Hands-On LLM Serving and Optimization

Hosting LLMs at Scale

With Early Release ebooks, you get books in their earliest form—the authors’ raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

Chi Wang and Peiheng Hu

Hands-On LLM Serving and Optimization

by Chi Wang and Peiheng Hu

Copyright © 2025 Chi wang and Peiheng Hu. All rights reserved.

Printed in the United States of America.

Published by O’Reilly Media, Inc., 1005 Gravenstein Highway North, Sebastopol, CA 95472.

O’Reilly books may be purchased for educational, business, or sales promotional use. Online editions are also available for most titles (http://oreilly.com). For more information, contact our corporate/institutional sales department: 800-998-9938 or corporate@oreilly.com.

Acquisitions Editor: Nicole Butterfield

Development Editor: Sarah Grey

Production Editor: Ashley Stussy

Interior Designer: David Futato

Cover Designer: Karen Montgomery

Illustrator: Kate Dullea

April 2026: First Edition

Revision History for the Early Release

2025-05-05: First Release

2025-06-10: Second Release

2025-08-08: Third Release

2025-11-11: Fourth Release

2025-12-17: Fifth Release

See http://oreilly.com/catalog/errata.csp?isbn=9798341621497 for release details.

The O’Reilly logo is a registered trademark of O’Reilly Media, Inc. Hands-On LLM Serving and Optimization, the cover image, and related trade dress are trademarks of O’Reilly Media, Inc.

The views expressed in this work are those of the authors and do not represent the publisher’s views. While the publisher and the authors have used good faith efforts to ensure that the information and instructions contained in this work are accurate, the publisher and the authors disclaim all responsibility for errors or omissions, including without limitation responsibility for damages resulting from the use of or reliance on this work. Use of the information and instructions contained in this work is at your own risk. If any code samples or other technology this work contains or describes is subject to open source licenses or the intellectual property rights of others, it is your responsibility to ensure that your use thereof complies with such licenses and/or rights.

979-8-341-62144-2

[FILL IN]

Brief Table of Contents (Not Yet Final)

Chapter 1: Introduction to Model Serving and Optimization (available)

Chapter 2: Large Language Model (LLM) Serving (available)

Chapter 3: Model Serving Best Practices and Case Studies (available)

Chapter 4: Build an Agent Application with LLM from Scratch (available)

Chapter 5: Performance Challenges When Serving LLMs (available)

Chapter 6: Common Optimization Techniques (available)

Chapter 7: LLM-Specific Optimization Techniques (available)

Chapter 8: LLM Serving Frameworks Overview (available)

Chapter 9: LLM Serving Optimization Best Practices and Case Study (unavailable)

Chapter 10: Efficiently Serve Multiple Fine-Tuned LLM Models (unavailable)

Chapter 1. Introduction to Model Serving and Optimization

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 1st chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

At its core, model serving is a process that addresses the challenge of making AI models accessible to end-users, applications, and systems, working through APIs, web services, or integrated workflows to generate predictions (called inferences) on new, unseen data.

To draw a simple analogy, to businesses of all kinds—whether they aim to deliver AI capabilities to their customers or enhance operational efficiency—model serving is a form of supply chain. For example, Amazon and Netflix use model serving to update customer recommendations instantly as users browse. Banks use model serving to block fraudulent transactions during online shopping checkouts, and airline chatbots use it to provide instant flight updates and rebooking options. Business stops for these companies if their model serving systems go down.

Just as every manufacturer strives for an efficient and cost-effective supply chain, AI companies seek to utilize their models and hardware efficiently and effectively in production environments. Therefore, choosing the right model serving approach and optimizing it are critical: it directly impacts the business’s lifeline and operational costs. Regardless of what role you have, if you are in the AI industry, you would benefit from having some knowledge of model serving.

As tech leaders who have worked on model serving infrastructure for over a decade, we have collaborated with a variety of participants in the field: researchers, developers, executives, marketing, customers, and students. I have found that people often feel intimidated or overwhelmed when they first attempt to understand model serving. There are three main reasons for this. First, you should already have a deep knowledge of model training. Second, as yet, there’s been no clear learning path that goes from get-started tutorials to managing a world-class serving system. Third, there are so many frameworks, libraries, vendors, and other engineering options that it can be hard to make decisions about which to adopt. This book aims to address the above challenges by providing a structured, practical guide to model serving and optimization—one that bridges the gap between theory and practice while empowering readers to make informed decisions.

In this chapter, we lay the foundation for you to understand the rest of the book. We begin by clarifying the core concepts of model serving, followed by a discussion of why robust, optimized model serving is critical for real-world applications. Finally, we explore general paradigms in model serving. By the end of this chapter, you will have a comprehensive overview of model serving and optimization, setting the stage for the hands-on materials in the subsequent chapters.

What Is a Model?

In academic terms, a machine-learning (ML) model is a mathematical representation or algorithm that learns patterns from data to make predictions, decisions, or inferences without being explicitly programmed for the task.

In engineering and operations, we focus more on how to use models than how to train them. Therefore we simply treat models in black-box fashion – as collections of (executable) files produced by ML training processes.

Figure 1-1. A model is composed of model architecture, model execution code and model data.

We view models as composed of three types of files: data, architecture, and execution code (see Figure 1-1):

Model data

A model’s data includes its weights, bias, and configuration. Weights and bias are what the model learns during training, and the model configuration holds the metadata to run the model, such as its embeddings and label classes (for classification models), its max_batch_size property (for batch inference), and its input and output tensors.

Model architecture

Architecture refers to the structure and design of an ML model. It defines how the model is organized, including the types and number of layers, the connections between layers, and the operations the model performs. The architecture determines how the model processes input data to produce output predictions or decisions.

Model execution code

A model’s execution code is what the model runs. It generally initializes the architecture in the model serving framework, loads weights, and runs predictions (or other outputs).

Tip

You can think of models as executable program files. Many people mistakenly assume that models are merely passive data files, but in reality, they also include execution logic, metadata, and other components that define how they operate. Rather than being static artifacts, models function as dynamic programs that can process inputs, make decisions, and evolve over time.

Let’s examine the structure of the model from Figure 1-1 using an example from the PyTorch tutorial – "Saving and Loading Models” – which saves a trained model to local drive.

Model Architecture

First, the following Python file defines the model architecture. This code (located in the model file/package) is normally a copy of the model architecture class in the model training code.

class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

Model Data

Second, after model training, the weights and biases are saved as a file called model_weights.pt.

// In PyTorch, the learnable parameters (i.e. weights and biases)
// are stored in "state_dict"
torch.save(model.state_dict(), model_weights.pt)

Model Execution Code

Third, the following example code defines how to execute the model for a given input:

model = TheModelClass(*args, **kwargs) //Initialize model
model.load_state_dict(torch.load(“model_weights.pt”, weights_only=True))
        //Load model weights (“model_weights.pt”)
model.eval() // set model to evaluation mode to run prediction
pred = model(inputs) // run model prediction with given input

Although different machine learning (ML) frameworks, such as TensorFlow and PyTorch, provide distinct APIs and libraries, the fundamental principles of model packaging and structure (Figure 1-1) remain largely similar. In practice, additional optimizations are often applied to model files to enhance performance for model serving, particularly for large language models (LLMs). These optimizations will be discussed in Chapters 6 and 7.

We recommend storing your model’s architecture and data in separate files

While it is possible to save an entire model in a single file (for example, torch.save(model, "model.pt")), separating the class definition (architecture) from the weights provides greater flexibility. This approach allows for more complex serving scenarios, such as loading pretrained weights into an updated model architecture while ignoring nonmatching keys or partially loading a model.

Now that you understand what an ML model is, let’s move on to model serving.

What Is Model Serving?

In engineering terminology, model serving refers to deploying an ML model in a production environment, where it can process new data and generate predictions.¹ This involves setting up the necessary infrastructure—both software and hardware—to ensure that the model can receive input, execute inference, and return results efficiently, scalably, and reliably.

Model serving can happen locally (on-device), in-cluster (on-premises), or remotely (on-cloud), depending on the business requirements and infrastructure constraints. Here are three examples, one for each serving scenario:

Warehouse robot (on-device)

A real-time object detection model (such as YOLO) is deployed on a robot’s onboard computer (say, a Raspberry Pi with a Coral TPU). This lets the robot process visual input and make autonomous decisions instantly.

Document search (on-premises)

To build a semantic search index on internal company data, a language model (likeSentence-BERT) is served within the company’s computing cluster. This setup generates dense vector embeddings for documents while ensuring data security and compliance by avoiding large-scale data transfers.

Customer support chatbot (on-cloud)

When a customer submits a query, the chatbot sends it, along with relevant documents, to a remote LLM hosted or provided by a vendor, such as Amazon SageMaker Inference or OpenAI, as shown in Figure 1-2.

Figure 1-2. Serving a LLM model for a chatbot app in SageMaker Inference Endpoint.

In the model-serving context, ML and DevOps engineers typically focus less on models’ internal details—such as their architecture, training methods, and file formats—and instead treat models as black boxes. This is because modern model-serving frameworks like vLLM and TensorRT-LLM handle the complexities of model execution for us, providing a high-level abstraction that makes running model predictions simpler and more efficient.

In model serving, we focus more on the following:

Deployment

Choosing the right hardware and making the model available from the training pipeline (or from open source options) to consume input data and return predictions.

Scalability and Availability

The ability to handle a few thousand to millions of requests efficiently while ensuring a consistent customer experience.

Latency

Delivering predictions quickly–in milliseconds, for real-time use cases.

Monitoring

Tracking model performance, data drift, and system health.

Versioning

Managing model updates and rollbacks without disrupting clients.

Security

Sensitive data must be protected, and access to the model should be controlled.

Cost to Serve

The most decisive factor of them all. We use cost-to-serve as a key factor to evaluate different serving approaches and trade-offs.

Model serving is a highly practical, engineering-focused field. Unlike AI research and model training, it does not require a deep understanding of ML algorithms or an academic background in AI. Instead, it emphasizes deploying and integrating models into real-world applications using existing tools and frameworks.

Note

For effective LLM serving and optimization, a solid understanding of AI algorithms, particularly Transformers, is essential. All advancements in LLM serving techniques are designed to overcome bottlenecks in LLM execution. So a strong grasp of LLM architecture provides us the intuition needed to optimize throughput and latency. In the following chapters, we will cover the foundational knowledge to help you understand the LLM serving and optimization methods.

Why Study Model Serving?

Most people would agree that model serving is critical, but aren’t sure why they need to learn to build or customize their own model-serving systems. In this section, we address some of these questions with our own experiences.

Question 1: “We consume models directly from cloud vendors such as Amazon AWS (SageMaker or BedRock) and Microsoft Azure ML. Why would we need to learn to build a model serving system?”

Deploying and hosting models on cloud vendors’ platforms is a good option in many situations. Cloud services can take care of many things, such as hardware management and software patching, and enable quick proofs of concept (POCs). We have successfully designed and built lots of model serving solutions on AWS and Azure . However, you can never simply consume these services without careful consideration.

While cloud-based model-serving solutions are convenient, they also come with challenges. First, there are numerous products and features to choose from, making it difficult to determine the best fit. Second, integrating vendor services with existing infrastructure and client applications often requires additional engineering efforts. Third, you’ll often need to optimize and fine-tune vendor services to align with your specific business use cases and performance requirements.

Cloud vendors design their model-serving solutions based on various customer needs and package them as general-purpose offerings. For instance, AWS SageMaker provides multiple serving options, including single-model and multi-model endpoints, and each option comes with different levels of customization, trade-offs, and costs. Without a clear understanding of model serving design, challenges, and use cases, it’s easy to make suboptimal choices that may increase costs or limit flexibility.

In short, understanding how a model serving system is built and operated helps you make the best use of cloud vendors’ model serving options.

Question 2: “Foundation LLMs are so powerful that we don’t need to consume any other models. We just let our client application call foundation model vendors, such as Open AI, DeepSeek, and Anthropic. Why would we need to build and maintain a model serving stack?”

LLMs are one of the greatest inventions in the AI field and a milestone on the path toward artificial general intelligence (AGI). But there is a lot of nuance when you actually apply it in business.

First, there are cost considerations. Businesses hardly ever use a single LLM to solve all of their business use cases. In the real world, businesses prefer to use smaller and cheaper models as much as possible and leverage LLMs for complicated requests. Also, when usage picks up, they generally want to try adopting lower-cost LLM serving. For example, running an open source LLM (such as DeepSeek) on an AWS Elastic Compute Cloud (EC2) instance (that is, a virtual server) can be 30 to 60% cheaper than doing so on a managed serving solution like AWS SageMaker. Also, you can often get a good deal (say, a 70% discount) if you reserve an EC2 instance, which further reduces the serving cost.

A typical chatbot app might use a smaller intent-classification model to triage customer questions, an embedding model to encode those customer questions and search for potential answers, and an LLM to break the customer’s request down into actions (backend API calls) and generate human-friendly user interactions.

Another aspect is data privacy and security. If your company is handling sensitive or proprietary data, you might be forced to leverage an in-house solution to avoid data breaches or compliance-related issues.

Lastly, most foundation-model vendors either don’t provide the functionality to serve your fine-tuned model, or charge a much higher price if you need to serve your own fine-tuned version of the model. For example, as we write this book in early 2025, OpenAI charges 50% more to serve a fine-tuned model. At the same time, smaller models are becoming much more powerful and providing fine-tuning for custom data and requirements, faster inference, better and more consistent results, and a lower cost-to-serve.

As technology keeps developing, it trends toward affordable and open approaches, and LLM businesses are no exception. There is nothing wrong with consuming LLMs from vendors, but having your own serving solution (on cloud or on-premises) with an open source LLM will give your business a critical edge at a lower cost.

Question 3: “Building and maintaining our own serving stack is expensive. What would make it worth the cost?”

Reducing model serving costs is an everlasting effort. As the industry develops and technology and infrastructure change, no one can guarantee that simply outsourcing your serving component will always meet your business’s goals. Understanding model serving deeply can help you understand all the trade-offs in order to make optimal choices. Let’s look at a few examples.

First, imagine you run a Series A startup in the education sector. Your app leverages LLMs to grade students’ homework and generate personalized study plans. Working with a fully managed LLM provider like OpenAI, DeepSeek, or Anthropic, is a great starting point to accelerate development and validate your business model.

However, as your usage scales, the serving costs start to become unsustainable. To optimize expenses, you consider migrating to a customized model serving solution: for example, using open-source LLMs, selecting optimized serving frameworks, and fine-tuning inference performance. This should significantly reduce your costs while maintaining flexibility and keeping you in control.

Let’s look at another scenario. This time, you run a customer relationship management (CRM) company that leverages LLMs to provide sales recommendations by processing data from multiple sources, including text-based inputs (emails, Slack messages, customer notes), unstructured audio searches (voice calls, meeting transcriptions) and structured CRM databases (past sales records, customer interactions).

Given the high volume of LLM inference requests and the involvement of sensitive customer data, fully outsourcing your model serving could introduce data-security risks and escalate your operational costs. You decide that owning and optimizing an in-house serving stack is necessary to lower the cost per inference, comply with security regulations (like GDPR and HIPAA), and achieve better performance and latency control.

Our experience has been that while outsourced solutions can be a great starting point, businesses must continuously assess whether in-house serving provides a long-term competitive advantage in terms of cost savings, security, and flexibility.

What We’ve Learned about Model Serving

The authors of this book have been building ML systems for a variety of enterprise use cases for over a decade. We’ve worked on projects that are fully outsourced (OpenAI), self-managed on-premise solutions (Tensorflow serving, TorchServe, Nvidia Triton, Ray), fully-managed cloud vendor solutions (AWS Bedrock), and deeply customized vendor solutions (AWS SageMaker).

What we’ve learned over the years is there is no “one size fits all” or “eternal” architecture. Model algorithms and serving technologies keep evolving, and businesses need to keep pace. Adopting the new technology lets you provide a better user experience and reduce operational cost, helping to keep the business alive and help it thrive.

As a business owner or model serving developer, it’s mandatory to have a solid understanding of the fundamentals of model serving, so that you can make sense of new innovations in the industry and assess pros and cons to make the right technical choices.

At the end of day, we don’t want you to be locked into any specific frameworks or vendors: we want you to have the ability to understand and adopt new serving technologies quickly, in order to give your business an edge among fierce competition.

Why Optimize Model Serving (Especially for LLMs)?

Once you’ve successfully deployed your model serving system, the next challenge is cost optimization. This raises key questions: Can we run the model on more cost-effective hardware? Can we improve throughput and reduce latency without upgrading infrastructure? Can we fully utilize the Cloud vendor’s LLM service? This is where model serving optimization comes into play.

Model serving optimization refers to the process for improving the model serving performance, such as reducing serving latency, increasing throughput, and optimizing resource usage. In practice, we do serving optimization to maximize serving efficiency while keeping costs under control.

Optimization Is Critical for LLMs

LLMs are large and complex enough to demand significant computing power, which can make their operational costs unaffordable for small businesses. For instance, Alphabet chairman John Hennessy told Reuters in 2023 that “running an LLM request can be 10 times more expensive than a traditional keyword search, potentially leading to billions in additional costs.”

So, for LLMs, it’s almost a must to optimize your model serving: doing so lets businesses reduce costs, enhance user experience with faster response times, and gain a competitive edge over companies with slower or more expensive inference solutions.

In the next section, we’ll break down the key concepts of model optimization and walk through some examples to help you get a feel for why an optimized model and serving system—especially for LLMs—can make a huge difference compared to a “raw” model straight from training.

No worries if some of this feels unfamiliar! We’ve got you covered—Chapters 5 through 8 will dive deeper into these concepts with hands-on examples to make everything clear. So stay tuned!

Why does model serving optimization work?

To get a good grasp of model serving optimization, it helps to first understand the difference between model training and model serving. A common question people ask is: “Why not just use the same setup for serving as we do for training? After all, both involve running the same model, right?”

Technically, that’s partially true—both processes involve loading the model and executing its architecture. But the goals and requirements are completely different. Let’s break it down in Table 1-1 to see why.

Table 1-1. Fundamental differences between model training and serving

Aspect
Model Training
Model Serving

Stage in ML Lifecycle

Prepares the model

Deploy to production

Objective

Run the model as a process to learn parameters (weights) by minimizing a loss function on the training data.

Run the model to generate prediction (inference) on new input efficiently.

Computation

Extremely compute-intensive, involving iterative model weight updates (including backpropagation and gradient updates)

Focused on efficient forward propagation only (no backpropagation)

Throughput and Latency

Preferable for high throughput (many samples per second). Large data batches are usually processed in parallel to optimize GPU utilization.

Optimized for low-latency response per request; often operates on single or small batches

Resource Requirement

Requires powerful GPUs/TPUs and distributed training frameworks (such as Fully Sharded Data Parallel (FSDP) and DeepSpeed)

Optimized for low-latency, resource-efficient execution, often on CPUs, edge devices, or inference-specific GPUs (such as NVIDIA TensorRT and ONNX Runtime)

From Table 1-1, you can see why using a model-training framework and system for serving would be inefficient. Instead, adopt model-serving-specific frameworks, such as NVIDIA Triton Inference Server, vLLM (Virtual LLM), and SGLang.

To give you an example, let’s look at vLLM, an optimized LLM serving engine designed by UC Berkeley’s SkyLab for high-throughput, low-latency inference. It significantly improves the efficiency of serving LLMs, with 24x higher throughput than basic serving approaches like Hugging Face Transformers or PyTorch’s native model execution, which are tailored to model training, static batching, being generic and user friendliness. Online serving, being more specific, poses more opportunities for improvement.

Like other model serving frameworks, vLLM exposes lots of knobs for engineers to tune for optimal model performance. Take a look at this vLLM command for running LLaMA-2 13B on an A100 80GB GPU, with a few optimal configurations:

python -m vllm.entrypoints.openai.api_server \
  --model llama-2-13b-chat-hf \
  --dtype bf16 \
  --gpu-memory-utilization 0.9 \
  --max-num-seqs 16
  --max-num-batched-tokens 16384 \
  --tensor-parallel-size 2

This vLLM command passes several configurations to vLLM aimed at achieving good throughput and latency on the Nvidia A100 GPU. For example, it sets Paged Attention (which is enabled by default in vLLM) to KV Cache, sets model precision to BF16, sets the max concurrent request (batch size) to 16 and shares the model across two A100 GPUs. (We’ll talk about all of these options in Chapters 6 and 7; it’s OK if they are unfamiliar now.)

Your Serving Framework’s Configuration Can Significantly Impact Performance

Every option available in a serving framework represents a potential optimization technique, and the right configurations can make a huge difference in performance.

For example, in our own experiment hosting the DeepSeek R1 model with vLLM, we gained a 15x increase in throughput—jumping from 38 tokens per second to 600 tokens per second—by doing two simple things in vLLM’s configuration: enabling FP8 Matrix Learning Accelerator (MLA) kernels and increasing the batch size.

Example: Performance Gain by Optimizing LLM’s KV Cache

We’ve covered a lot of high-level optimization concepts so far, but now it’s time for you to see their impact in action.

In this section, we’ll share some performance improvements by optimizing the key-value (KV) cache for LLMs using two serving frameworks—vLLM and SGLang. These results will give you a tangible sense of how model optimization translates into real-world improvements.

What is LLM (Transformer) KV (Key-Value) Cache?

Key-value caching is a technique that helps speed up the model inference process by remembering important information (such as attention) from previous steps. Instead of recomputing everything from scratch, the model reuses what it has already calculated, making text generation much faster and more efficient. Figure 1-3 shows a general KV-cache workflow in the transformer decoding process.

Figure 1-3. How a KV cache works, step by step. (Source: HuggingFace blog )

In Figure 1-3, you can see how the KV cache saves unnecessary attention calculations:

When the model sees the input prompt, it calculates and stores it as key-value pairs in the cache.

When generating new tokens, instead of starting over to recalculate KV cache from the very beginning, the model retrieves the stored KV cache instead.

With the KV cache, the model now can calculate attention efficiently by aligning the cached keys and values with the new query (Q) to compute the new token.

The model appends the newly generated token to the existing sequence and repeats the process from step 2 until it is finished generating.

The KV cache makes a big difference to the speed and efficiency of LLM inference, especially for long context and long generations. By saving and reusing past calculations, it avoids the need to start over each time and trades storage space for less computation, making it much faster than the regular way of generating text. Even a vanilla version of a KV cache can lead to a fivefold improvement in LLM serving speed, as the HuggingFace (HF) blog explains.

Now that you know what a KV cache is, let’s look at what happens when we use vLLM and SGLang to optimize an LLM’s KV cache.

vLLM with Paged Attention

Increasing an LLM’s throughput requires batching many requests at a time. However, the KV cache memory for each request is huge, and it grows and shrinks dynamically. This means that, if it is managed inefficiently, this memory can be significantly wasted by fragmentation and redundant duplication, limiting the batch size and thus affecting the throughput.

vLLM comes with a good approach, a feature called PagedAttention that manages the KV cache dynamically, like a paging system in an OS, to eliminate fragmentation and enable continuous token generation across multiple requests without memory waste. As a result, research has shown that it improves GPU utilization.

The vLLM team experimented in 2023 by adopting PagedAttention on two settings: LLaMA-7B on an NVIDIA A10G GPU and LLaMA-13B on an NVIDIA A100 GPU (40GB). Their results were impressive, achieving throughput up to 24x higher than HF Transformers and up to 3.5x higher than HF Text Generation Inference (TGI). Figure 1-4 illustrates their findings.

Figure 1-4. vLLM achieves 8.5x - 15x higher throughput than HF and 3.3x - 3.5x higher throughput than TGI. Image source: vLLM blog

You can use gpu-memory-utilization, max-model-len, and block-size to fine-tune the PagedAttention feature based on model size:

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-2-13b-chat-hf \
  --gpu-memory-utilization 0.9 \
  --max-model-len 4096 \
  --block-size 32

More details of vLLM configuration can be found in the vLLM documentation.

SGLang with RadixAttention

SGLang is a model-serving framework that focuses on optimizing the KV cache across multiple LLM requests. Since KV cache computation depends on prefix tokens, requests with the same prompt prefix can reuse the KV cache, reducing redundant computation and memory usage.

In existing inference engines, the KV cache for a request is discarded after its own generation is completed and is not reused across multiple calls. If the input prompts of the requests happen to have many common prefixes, recomputing all the KV cache is not efficient and reusing them can vastly reduce the serving latency. So SGLang came up with a feature called RadixAttention, which maintains a LRU cache of the KV cache for all requests within a radix tree; it enables automatic reuse of the KV cache across multiple generation calls efficiently.

As a result, in an experiment on hosting open-weight models (Llama-7B and Mixtral-8x7B models), researchers found that SGLang improves throughput by up to 6.4x and reduces latency by up to 3.7x.² See more at paper, SGLang: Efficient Execution of Structured Language Model Programs.

I hope these examples give you a clear intuition on how optimizing the KV cache computation can significantly boost model serving performance.

More importantly, I want to emphasize just how significant an impact model optimization can have—especially on LLMs. Given the high cost of GPU-based inference, achieving 6x throughput and 3x lower latency without any additional infrastructure cost is a game-changer.

Model Serving Paradigms

In this section, we’ll explore the most common model serving paradigms and their real-world applications. We’ll introduce these designs progressively—starting with simple on-device and single-model serving, then expanding to more complex multimodel architectures, each building upon the previous concept.

Along the way, we’ll share insights from our real-world experience to help you understand the trade-offs of each approach. By the end of this section, you’ll have a clear sense of which serving solution best fits your business needs and a foundational understanding of high-level design paradigms for structuring it. The practical design principles, along with hands-on examples of model serving systems, will be covered in Chapters 2, 3, and 4.

On-Device (Edge) Serving

On-device serving refers to running the model directly on user-side devices, rather than relying on remote web services. This setup allows applications to process data locally—where it is generated—on devices such as smartphones, drones, robots, cameras, and VR headsets, enabling real-time computation without the need for constant internet connectivity. Figure 1-5 highlights the four key benefits of on-device AI.

Figure 1-5. The benefits of on-device model serving (Image source: DeepLearning.AI)

Given these benefits, on-device model serving has been widely adopted across various applications. Examples include smartphone face ID, AI-powered background-noise reduction in calls, intelligent robotic vacuum cleaners, and drones equipped with AI-assisted flight mode, all leveraging local AI processing for faster, more efficient performance.

On-Device Serving Design

Now let’s explore how on-device model serving operates. We’ll begin by examining the key components of an on-device AI application, followed by an explanation of the deployment process required to prepare a model for different devices. The core concepts are illustrated in Figure 1-6.

Figure 1-6. On-device AI app design and model deployment workflow.

In the high-level n-device AI app design shown in Figure 1-6 (a), model runtime and model wrapper are the key components.

A model runtime is a specialized software framework designed to execute ML models efficiently on different hardware platforms. It serves as the intermediary layer between the trained model and the device’s hardware, optimizing the inference process to maximize speed, efficiency, and resource utilization.

The model runtime plays a crucial role in on-device AI by abstracting the complexities of model execution across different hardware (like smartphones, drones, and robots) and operating systems (iOS, Android, Linux). It also provides hardware-specific optimizations and acceleration techniques to improve inference performance.

For example, the model runtime can leverage specialized hardware (like CPUs, GPUs, NPUs, and TPUs) for efficient inference. It also supports delegates (such as GPU Delegate, NNAPI for Android, and Core ML for iOS) to optimize performance.

A well-designed model runtime allows developers to focus on building application features rather than worrying about hardware compatibility or model execution details. As of this writing, some of the most popular model runtimes include TensorFlow Lite (TFLite), ONNX Runtime (ORT) and Core ML.

As AI app developers, we use model runtimes to execute models efficiently and encapsulate model-specific execution logic into a reusable module often called a model wrapper. This wrapper abstracts low-level details, making it easier to integrate, update, and optimize models within applications.

A model wrapper is a component implemented by the application developers to provide the model inference interfaces (functions) for the application logic to execute the model. The model wrapper encapsulates the details of how to interact with the model runtime, such as preprocessing the input data to model input, load the model into the model runtime, execute the model from the model runtime, and postprocess the output.

In the workflow shown in Figure 1-6 (a), the application logic first asks the model wrapper to execute a model with given data. The model wrapper takes care of data conversion and calls the model runtime to execute the model. The model runtime then handles the execution using the local hardware.

Preparing a model for on-device serving

Figure 1-6 (b) describes the high-level workflow of getting a model ready, from the training process to operating on-device.

Once the model has been trained, you first need to convert it from its training format into the format the model runtime requires. For example, if you’re using the Tensorflow Lite runtime, you can convert (compile) a Resnet pytorch model into Tensorflow Lite’s file format (.tflite) by using the following command:

ai_edge_torch.convert(resnet18.eval(), sample_input)

Once the model is compiled, the next step is to validate its numerical accuracy on the target device. This ensures that converting the model and optimizing it for on-device execution haven’t degraded its accuracy. The most common validation method is to run the same model inputs on both the training server and the local device, then compare the output differences.

Next, you’ll need to measure the model’s performance to ensure that the local hardware can execute the model efficiently while meeting the business requirements. Once you’ve validated that, you can package the optimized model as part of the application installation or update for release.

This deployment process may seem tedious, but tools like Qualcomm AI Hub help automate all the key steps, including on-device model compilation, optimization, and validation, making the workflow much more efficient.

Benefits, Limitations, and Challenges

Running models directly on devices offers benefits like low latency, offline capabilities, and better privacy, but it comes with several limitations and challenges and often requires cloud-based model serving. Here’s a breakdown of the key challenges, along with examples.

Computation and Storage Constraints

Smartphones, IoT hardware, and embedded systems have limited CPU, GPU, and memory resources, making it difficult to run large, complex models. For example, running a large transformer-based language model like GPT on a smartphone is infeasible due to memory, storage, and processing-power constraints.

Power Consumption

Running AI models, especially deep neural networks, can be power-intensive, draining battery life quickly. For example, we don’t run real-time video enhancement (from low to high resolution) locally, since it consumes excessive amounts of power.

Limited Update and Maintenance Capabilities

On-device models are hard to update frequently. When improvements are made to the model, each device must receive and install an update. For example, a banking app using an on-device fraud-detection model would need to release frequent app updates to improve its model’s accuracy.

Inconsistent Hardware Support

Different devices have different hardware capabilities: some support NPUs, others only CPUs. This can make it hard to deploy optimized models across all devices. For example, a model optimized for Apple’s Neural Engine on an iPhone might not run efficiently on an Android device with a Qualcomm Snapdragon processor.

When faced with these challenges, we often choose to either move the entire model serving to the cloud or adopt a hybrid approach, where lightweight AI processing happens on-device and more complex computations are offloaded to the cloud. Next, let’s take a closer look at cloud-based model serving approaches.

Single Model Service

In this section, I’ll introduce you to the classic Single-Model Service design pattern, which is also the most widely used cloud-based model-serving approach.

In this design (diagrammed in Figure 1-7), each model, model version, or model type is deployed as a dedicated web service, exposing a prediction API over HTTP or gRPC for server-side execution. This setup ensures scalability, isolation, and flexibility in model deployment.

Figure 1-7. Single-model service architecture.

From Figure 1-7, you can see that the Single-Model Service architecture follows a standard microservices design. It has an API component in front to receive web prediction requests and route them to one of its backend workers/containers. The backend workers and containers perform the actual model execution and return the result.

Containerization is the Foundation of Modern Model Serving

Containerization means creating an isolated environment (or container) that packages an application along with all its dependencies, libraries, and configurations. Using a container ensures that the application runs consistently across different computing environments, whether on a developer’s laptop, a test server, or in production.

In modern model serving, most server-side model execution happens inside a container. It’s a common practice to encapsulate the model management, resource management and model execution logic into a Docker container (or Kubernetes Pod), and expose the model’s prediction/serving function via a gRPC or HTTP interface.

Docker is the most popular containerization technology as we write this book, so we use Docker containers throughout the book, including code examples.

Now let’s look at Figure 1-8 to see what’s inside a single model serving container.

Figure 1-8. Single-model serving container design

You can see that there are three main components in the model-serving container (Figure 1-8): API-Server, Model Management, and Inference Backend. Inference Backend is the model execution component. It usually leverages a model-serving framework (like TensorFlow Serving or TorchServe) or libraries (like vLLM or TensorRT-LLM) to run model serving. Model Management is the component responsible for preparing the model for serving. This includes downloading the model, extracting it to local storage, and loading it into the inference backend ( initialization steps a through d in Figure 1-8). API-Server exposes the model inference functionality over HTTP or gRPC, enabling external applications to send prediction requests.

For model deployment, we expect the model-training pipeline or the Ops (Operations) team to upload the trained model to cloud storage. The Model Management component then automatically refreshes the model in each serving container, ensuring that the latest version is deployed seamlessly.

Routing Choices in Single Model Serving

One of the key responsibilities of the model service API component (in Figure 1-7) is to ensure all the model serving containers are saturated with the prediction workloads equally. It sounds like a problem that can be solved easily using a “round-robin” load-balancing strategy. For example, if you have three containers (c1, c2, and c3), the first prediction request would go to c1, the second to c2, the third to c3, and then it starts over again.

However, round-robin routing has a key limitation—it doesn’t account for varying request processing times. In many cases, prediction requests have different computational costs based on input size and complexity. For example, an LLM inference request with a large text payload takes longer to process than one with a shorter input. An image model processing high-resolution images consumes more GPU/CPU resources than one processing smaller images.

Round-robin load balancing creates inefficiencies because of this variation: for instance, you might have a situation where container c1 is still processing a large request and container c2 has already finished its work, but the next request is still routed to c1, creating a bottleneck.

To address this issue, we need more sophisticated routing strategies that take into account processing time, resource availability, and request complexity. Here are some frequently used routing choices:

Weighted Round Robin

Servers (bigger) with higher weights receive more requests

Least Connections

Directs requests to the server with the fewest active connections

Least Response Time

Routes requests to the server with the lowest response time

Dynamic Load Balancing

Uses metrics like CPU, GPU, memory usage, or queue length to distribute requests

Horizontal and Vertical Scaling

As prediction-request traffic increases, a single serving container may struggle to process requests quickly. So we need to deploy more instances of the serving containers (we’ll refer to them as just instances) across multiple machines to handle the growing number of concurrent users with minimal delays. This is called horizontal scaling, or scaling out.

In a production serving environment, we enable autoscaling, so that the infrastructure management system automatically provisions more instances based on server resource utilization. The system monitors key metrics such as CPU and memory usage, inference latency, and number of active requests.

When monitored metrics exceed their predefined thresholds, the infrastructure management system spins up additional instances dynamically. For example, if you use Kubernetes to manage your compute cluster, its Horizontal Pod Autoscaling (HPA) can automatically adjust the number of serving pods based on real-time demand.

For large deep-learning models, such as GPT-4 and Llama 2-70B, a single machine may not have enough GPU memory to load the model. In these cases, we use vertical scaling (scaling up). This means deploying the model on more powerful GPUs (such as NVIDIA H100 or A100) and using distributed serving across multiple GPUs or multiple machines to split the workload.

Distributed model serving can be complex, but modern frameworks like vLLM simplify the process by abstracting away the low-level details. You can easily specify the number of GPUs to use when launching a model. For example, to serve a Llama model on 4 GPUs using vLLM, you can run this command:

python -m vllm.entrypoints.api_server \
   --model meta-llama/Llama-2-13b-hf \
   --tensor-parallel-size 4 \

Prioritizing Intra-Node Serving (Multiple GPUs on One Machine) Over Inter-Node Serving (Multiple Machines)

For better serving performance, easier management, and lower latency, it’s generally preferable to run large models on a single machine with multiple GPUs (intra-node serving) rather than distributing them across multiple machines (inter-node serving).

Inter-node serving introduces network overhead, increasing latency and complexity in synchronization. Whenever possible, we recommend you optimize model execution to fit within a single machine, to maximize efficiency.

We’ll dive deeper into these model optimization techniques in Chapters 6 and 7.

Benefits, Limitations, and Challenges

The Single Model Service approach is our favorite because of its simplicity and versatility. It works seamlessly with any type of model, small or large, from LLMs to deep learning models, making it a reliable and scalable choice for model serving.

From a production and operation perspective, we use Single Model Service whenever possible for the following reasons:

It offers the best performance, with no resource contention and lower latency.

Scaling is easy, since each model can scale independently.

Deployment and debugging are simpler thanks to isolated logs, metrics, and updates.

It’s more reliable since if one model crashes, it doesn’t affect others.

Optimized hardware per model allows for better cost control.

Although the Single Model Serving approach is great, it falls short, particularly in resource efficiency and cost.

Imagine you’re running an agent platform where customers can create their own agents and deploy custom models. If you have 100 customers and each deploys 10 models, that means your platform would be hosting 1,000 separate model services. This setup isn’t scalable—the overhead of maintaining, patching, deploying, and monitoring so many services would be overwhelming. Worse, many uploaded models might never be used, leading to wasted compute resources and unnecessary infrastructure costs.

To reduce operational complexity and maximize resource utilization, you need a more efficient solution—one that allows multiple models to share compute resources dynamically, instead of running a separate service for each.

This brings us to Multi-Model Service, a design approach that enables cost-effective, scalable model serving.

Multi-Model Service

Multi-Model Service is designed to co-host multiple models in one serving container, sharing GPU/CPU and memory across multiple models, and to load and unload models dynamically based on the incoming traffic. This can help you save significantly on costs and achieve the best price performance.

Multi-Model Service is an ideal solution when you need to serve a large number of models efficiently without loading them all into memory at once. Instead of dedicating a separate container to each model, this approach has models share compute resources within a single serving instance.

Going back to our agent-platform use case, instead of loading all 1,000 models into GPU and memory simultaneously, with Multi-Model Service, you could adopt an on-demand approach: Load the model only when a customer requests a prediction, and unload it when it’s inactive or when memory needs to be freed up for a new model. This significantly reduces infrastructure costs and improves resource utilization, making the model serving far more scalable.

You might wonder how it’s possible to host multiple models in the same serving container when they may have been trained using different frameworks and require different execution parameters?

The answer lies in the design of the Model Serving container, which manages model loading, execution, and unloading dynamically. Let’s explore its architecture, shown in Figure 1-9.

Figure 1-9. Multi-Model Serving Container Design

If you compare Figure 1-8 and Figure 1-9– the Single Model and Multi-Model serving containers, you can see that the Multi-Model Serving Container has two components: a Model Server Inference Backend and Model Cache Management.

The Model Server Inference Backend is designed to handle multiple types of models in a black-box manner. Regardless of the model algorithm and version, it can serve these models with an unified web prediction API. The trick here is that the model server contains multiple types of inference backends internally, and each backend handles one type of model (for example, TensorFlow, ONNX, or PyTorch).

Fortunately, you don’t have to build a model server from scratch—there are robust open-source solutions available, such as NVIDIA Triton Inference Server. In practice, we deploy systems like Triton inside model serving containers (or Kubernetes pods), using it as the Model Server Inference Backend to efficiently host and manage multiple models in GPU and memory simultaneously.

Note

Triton Inference Server is one of the most popular open-source model serving solutions, and it’s also one of our go-to tools, especially for multi-model use cases.

Triton can manage and serve any number or mix of models, constrained only by system disk and memory resources. It provides a unified API on top of various model formats, including: TensorRT, TensorFlow GraphDef, TensorFlow SavedModel, ONNX, PyTorch, Caffe2 NetDef and more. See more about Triton from Triton Architecture Blog.

The Model Cache Management component has two responsibilities. First, it loads a given model from model storage into the Model Server Inference Backend. Second, it prevents the container resources from model overload by maintaining a Least Recently Used (LRU) cache to track all the modeled models, so it can unload the least-used models when the container’s resource utilization is high.

Now, let’s look at an end-to-end on-demand model serving workflow. When a prediction request on model A hits the container, model management will do three things for the request:

If model A is loaded already, it forwards the request to the model server backend to execute the model and return the result.

If model A is not loaded, the management component will download the model from model storage, load it into the model server, and forward the request.

If memory and disk consumption is above a specified threshold (for example, 80%), the management component will find the least used model in its cache, unload it from the server backend, and remove the model from cache and disk.

Routing and autoscaling in Multi-Model Service

Sharing resources to host multiple models is efficient and cost-effective, but it also introduces challenges in routing prediction requests to the right models and scaling them effectively. Let’s dive into these challenges and explore potential solutions.

Figure 1-10. Serving and scaling challenges in multi-model serving

We can see two problems in Figure 1-10.

The first problem is about routing. Since models can be hosted in any of the serving containers, we want to route the prediction request to the container which already loaded the model. Otherwise, we add extra latency to the request issue: “cold start” (pending the request for model loading) and “model swapping” (removing old models out of memory for the new requested model).

The second problem has to do with per-model scaling. Not all models will take the same amount of traffic; some models will be used more frequently than others. We want to have more model instances in our multi-model containers for these “hot” models.

To solve the routing and model-scaling issues, we need to enhance the routing component at the frontend API layer (see Figure 1-11).

Figure 1-11. Route model server requests with host/model map and model replica counts

In Figure 1-11, we add two features to address the model-routing and autoscaling problems that arise in multi-model serving. First, we add an attribute “replica” to the model’s metadata. The number of replicas defines how many instances of this model we want to host in the backend containers. Second, we introduce a route map in the routing layer, so it knows to which containers it should route the request.

For example, as Figure 1-11 shows, for a given Model A with a current replica count of 2, the routing logic will choose two containers (for example, c1 and c3), to host it. Then it will route all future model A requests to c1 and c3. Depending on Model A’s traffic ratio, the routing component can update its replica count to autoscale Model A up and down.

Benefits, Limitations, and Challenges

As we’ve explained, multi-model service is the most cost-efficient serving method since models can share the same compute resources, which is extremely powerful when you don’t have to host all the models at the same time.

However, multi-model service tends to struggle in the following situations:

When the models are too large to fit in one GPU. Either there is no room to share or lots of model unloading is required to make memory space available. This causes a lot of model loading overhead and cold-start latency for future models.

When individual models have high traffic levels, which requires low-latency, always-loaded models. Single-model service is better in these cases.

When the models have different security policies.

When operational complexity is high. For instance, deployment, debugging, and monitoring are complicated due to the number of moving pieces in the system, such as model cache management, per-model routing and scaling, model compatibility, and dependency conflicts.

At this point, you might already have noticed that the multi-model and single-model serving approaches are complementary. In practice, we compose different serving approaches together to make up a platform, so we can address different serving use cases under the same roof. The next section discusses model serving platforms.

Model Serving Platforms

As a business grows, its model-serving demand surges, and its requirements become increasingly complex. Simply putting single-model and multi-model services together is no longer good enough. These serving systems face two major challenges.

First, many tasks now require multiple models to work together. For instance, a voice assistant like Apple Siri must integrate speech recognition, natural language processing (NLP), recommendation models, and text-to-speech synthesis to respond to a single voice command. Managing these multi-model interactions efficiently is crucial for seamless performance.

Second, when a growing number of AI applications are running on one platform, optimizing compute resources becomes a critical challenge. You need an efficient strategy to allocate GPUs, CPUs, and memory across models while ensuring scalability, cost-effectiveness, and minimal resource contention.

Figure 1-12. Model Serving Platform Design

The model serving platform design in Figure 1-12 gives you a sense of how a model serving platform addresses the aforementioned challenges. First, it uses resource groups to segregate the prediction workload for different applications. For example, we see two single model services and one multi-model service deployed in the App1 resource group. Each resource group is allocated CPU, GPU, and memory quotas based on business requirements to optimize performance and cost.

Second, the design has a graph execution component (for example, Airflow or Ray) to support multistep inference workflows. Each app team will define their own inference workflow and deploy it to the graph execution engine. Then the graph engine runs the workflow. For each inference step in the workflow, the graph engine calls the routing component to find the right serving group and service to run inference.

Figure 1-12 provides a high-level overview for a model-serving platform. In practice, this design could be a lot more complicated. For example, a serving platform must have security and access-control components, metric and monitoring components, deployment and DevOps integration, and more. Since these components are not core to model serving, we’ve omitted them from this diagram to keep the focus on the main serving concepts.

Summary

In this chapter, we explored the structure of MLfrom an engineering perspective, viewing them as executable components rather than just static data files. A model consists of three key elements: model data (weights, biases, and configurations), model architecture (the structure of layers and operations of the model), and model execution code (which loads and runs the model for inference).

Next, we introduced model serving, which involves deploying models in production environments to process real-time data and generate predictions. We highlighted three deployment scenarios—on-device, on-premises, and cloud-based model serving—and outlined the key considerations for model serving, such as scalability, latency, monitoring, security, and cost optimization.

Then we explored why understanding and optimizing model serving is essential, even when leveraging cloud vendor solutions or large language models (LLMs). While cloud-based model serving is convenient, businesses must still navigate complex integrations, cost trade-offs, and security concerns. Similarly, LLMs are powerful but often require multi-model pipelines, fine-tuning capabilities, and cost-efficient deployment strategies to be practical in real-world applications.

We also introduced model serving optimization, particularly for LLMs, where efficiency gains can lead to significant cost reductions and performance improvements. By optimizing techniques such as KV Cache with Paged Attention (vLLM) and Radix Attention (SGLang), businesses can achieve six times higher throughput and a third of the latency without incurring additional infrastructure costs. Since LLM inference is computationally expensive, serving optimization is not just an improvement—it’s a necessity for cost-effective AI deployment.

Finally, we explored the most common model-serving paradigms, from on-device (edge) serving to single-model service, multi-model, and full-model serving platforms. Each paradigm has its strengths and trade-offs, depending on factors such as latency, cost, scalability, and complexity.

We hope this chapter has convinced you that model serving is critical–and that it’s an engineering-driven discipline, focusing on efficiently deploying and integrating models into real-world applications, rather than on the intricacies of training models or designing algorithms.

In the next two chapters, we’ll dive into hands-on examples to see how model serving is built and operates in practice. We’ll start with LLM model serving in Chapter 2, then expand to a broader model-serving system design and implementation in Chapter 3.

¹ We use the terms model serving, model inference, and prediction interchangeably throughout this book.

² The throughput and latency improvements SGLang made come not only from KV cache reuse, but also from exploiting parallelism within a single program and from faster constrained decoding.

Chapter 2. Large Language Model (LLM) Serving

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 2nd chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

In the previous chapter, we introduced the concept of models, model serving, and common serving paradigms. In this chapter, we shift our focus to the specific challenges and techniques involved in serving large language models (LLMs).

One of the most common barriers for those entering the world of model serving is the sheer complexity of modern serving systems. With rapidly evolving model architectures, training algorithms, and tooling—combined with layers of production infrastructure like monitoring, scaling, security, CI/CD pipelines, and service dependencies—the experience can quickly become overwhelming. The result is that many engineers and researchers lose focus, getting bogged down in system details before they can grasp the core principles.

To address this, our approach is to start from the fundamentals. We begin with the bare minimum code required to serve an LLM and build up from there. This helps establish a strong mental model of how token generation works and why serving LLMs poses unique challenges. From this foundation, we’ll gradually expand into more advanced topics like system architecture, performance optimization methods, and infrastructure choices in the rest of the chapters.

In this chapter, we start with:

The basic architecture of LLMs, including the token generation process and attention mechanisms

What happens under the hood during inference through hands-on code examples

The core concepts behind LLM serving, such as prefill, decode, and key-value (KV) cache reuse

Why understanding these fundamentals is critical to diagnosing bottlenecks and contributing to performance improvements

We’ll then transition into using a modern serving framework—vLLM—to demonstrate how to improve model serving efficiency. Next, using vLLM, we will introduce key serving methods such as streaming and batching, which are essential to achieving production-grade performance.

This is a foundational chapter in the book. Our goal is to equip you with the practical understanding and intuition needed to reason about LLM serving systems. You’ll gain clarity on their behavior, limitations, and optimization opportunities—laying the groundwork for later chapters, where we dive deeper into system design, scalability, and advanced serving optimization strategies.

Throughout the chapter, we provide hands-on examples to help reinforce each concept. And don’t worry—deep mathematical knowledge is not required. Since we are focused on serving rather than training, we’ve abstracted mathematical concepts into intuitive explanations, allowing you to focus on the engineering perspective.

In the next chapter, we’ll take what we’ve learned here and demonstrate how to wrap it into a real-world web service, covering key design decisions and principles.

Inside the Mind of a Transformer

In this section, we’ll introduce the essential concepts behind the Transformer model—specifically from a model serving perspective. Rather than diving into training algorithms or academic theory, we take a top-down, conceptual approach: beginning with the history of LLMs, then walking through the LLM generation process, model architecture, Transformer blocks, and finally the attention mechanism.

LLM Evolution history

The evolution of LLMs is not just a historical curiosity—it offers critical insights into the model design choices, architectural patterns, and execution behaviors which is fundamental to inference and optimization workflows.

Language models have evolved from basic rule-based systems to sophisticated neural networks capable of generating coherent and contextually relevant text. This progress has been driven by advancements in model architectures, the availability of large-scale text datasets, and increasing computational power. Figure 2-1 provides an overview of language model development.

Figure 2-1. The history and development of language models. (Source)

The late 2000s marked a turning point in natural language processing (NLP) with the advent of deep learning (for the sake of keeping this brief, we’ll skip over early rule-based methods and n-gram models here). A major breakthrough came in 2013 with Word2Vec by Mikolov et al. at Google, which introduced dense vector representations of words in a continuous space. These embeddings captured semantic relationships between words and dramatically advanced the field’s ability to understand language.

To effectively model the sequential nature of language, recurrent neural networks (RNNs) became popular in 2013. RNNs could capture temporal dependencies and contextual information, making them well-suited for tasks such as sentiment analysis and text generation. To overcome RNNs’ struggles with long-term dependencies, more sophisticated architectures like long short-term memory (LSTM) networks and gated recurrent units (GRUs) were introduced in 2014. These incorporated gating mechanisms and memory cells to better retain and manage information over longer sequences.

Despite these improvements, RNN-based models still had critical limitations. For instance, they processed input sequentially—one time step at a time—making them inherently difficult to parallelize. This constraint hindered their scalability and efficiency on modern hardware like GPUs, and they continued to struggle with capturing long-range dependencies in text.

The introduction of the Transformer architecture, introduced by Google in the 2017 paper "Attention Is All You Need," revolutionized sequence modeling. By replacing RNNs’ recurrent layers with self-attention mechanisms and positional encoding, Transformers enable parallel processing while efficiently capturing long-range dependencies.

The introduction of the Transformer architecture sparked the emergence of two influential model families: Bidirectional Encoder Representations from Transformers (BERT ) and Generative Pre-trained Transformer (GPT). A pivotal innovation was the shift from training models from scratch for each task to pre-training on large-scale, unlabeled text using unsupervised learning objectives. This pre-training phase enables models to acquire a broad, generalizable understanding of language, which can then be fine-tuned for specific downstream tasks.

BERT, built on a bidirectional encoder, reads text in both directions simultaneously. This makes it particularly effective for understanding context and performing tasks such as text classification or generating contextual embeddings. In contrast, GPT models are based on a unidirectional decoder, generating text by predicting the next token given the preceding context—an approach well-suited for generative tasks.

Over time, the GPT architecture has demonstrated exceptional versatility across a wide range of applications, including text generation, summarization, translation, and question answering. Researchers have discovered that scaling up both model size and training data significantly enhanced performance. This scaling unlocked capabilities like few-shot and even zero-shot learning, where models can perform new tasks with little or no additional training.

This trend is evident in the rapid evolution of model scale: GPT-1 had 117 million parameters, while GPT-3 grew to 175 billion within two years. More recently, DeepSeek R1 reached 671 billion parameters. As a result, the term large language model (LLM) emerged to describe these powerful pre-trained models characterized by tens or even hundreds of billions of parameters.

The definition of LLM will keep evolving

As of the time of writing, the term LLM typically refers to generative, decoder-only Transformer models. However, it’s important to understand that this definition is not static—it continues to evolve alongside advancements in AI research and deployment.

At its foundation, LLMs are built on the concept of large-scale language modeling. But what qualifies as “large” has changed dramatically over time, as have the architectures, capabilities, and training methodologies involved. For instance, in 2018, models like BERT—with a few hundred million parameters—were considered large. By 2023–2024, models such as GPT-4 and Claude 3 are estimated to scale to trillions of parameters and include multi-modal capabilities.

Alongside this growth, the expectations for LLMs have expanded significantly. They are no longer judged solely by their ability to generate coherent text, but also by their capacity to follow complex instructions, interact with tools, and exhibit human-aligned behavior.

Therefore, readers should remain mindful that both the definitions and the expectations surrounding LLMs will continue to evolve as the field progresses.

With the brief historical overview out of the way, we now turn to a more technical discussion of Transformer architecture. If you are interested in a deeper exploration of LLM history and development, the paper “A Survey of Large Language Models” offers a comprehensive and insightful read.

Note

In this chapter, the terms LLM and Transformer specifically refer to Decoder-Only Transformer architectures—models that utilize only the decoder stack from the original Transformer design introduced and popularized by OpenAI’s 2018 paper “Improving Language Understanding by Generative Pre‑Training.”

Please note, however, that we use the terms input text, input sequence, and prompt interchangeably throughout the book to mean input to an LLM.

The Autoregressive Nature of Transformer

In this section, we’ll take a high-level, conceptual look at how a Transformer performs text generation given an input prompt.¹ A defining characteristic of LLMs is their autoregressive nature: they generate text one token at a time, with each new token predicted based on all previously generated tokens. This step-by-step process allows the model to maintain coherence and context, ensuring that each word aligns meaningfully with what has already been generated. Autoregressive generation also mirrors how language is naturally constructed—progressively, with each word depending on prior context. See Figure 2-2 for an illustration of this process.

Figure 2-2. Transformers generate tokens one at a time. The output token will be appended to the input sequence as a prompt for the next token generation.

In Figure 2-2, you can see how a Transformer-based LLM generates text in an autoregressive manner—predicting one token at a time, each conditioned on the prompt and all previously generated tokens:

In step #1, the model receives the initial prompt, “Write a short introduction about the US capital city,” and generates the first output token: Washington.

In step #2, this generated token is appended to the original input, and the updated sequence is fed back into the model to produce the next token: D.C.

In step #3, the prompt now includes both Washington and D.C., allowing the model to generate the next token: is.

In step #4, the sequence continues with the addition of is, leading the model to generate the.

This iterative process continues until the model encounters a stopping condition—such as reaching a maximum length or generating a special stop token—gradually building the output sequence token by token.

Decoder-Only Transformer Architecture

Now that you’ve seen how a Transformer model generates text, let’s take a closer look inside the “Transformer LLM” box shown in Figure 2-2.

In this section, we’ll explore the underlying architecture of decoder-only Transformers, which power most modern LLMs. In addition to reviewing the high-level components and flow, we’ll also examine the architecture of a real-world model—Qwen 2.5—to help build your practical intuition for understanding and analyzing LLM structures in practice.

Why demonstrate decoder-only Transformer architecture?

There are many architectural variants built on the Transformer framework. While these models all share the same foundational principle—stacked Transformer blocks with self-attention—they differ in how their components are structured and applied.

In this book, we focus on the decoder-only Transformer architecture because it is the most widely adopted design for models that perform generative tasks, such as GPT, LLaMA, and Qwen. If you are interested in exploring other Transformer architectures, such as encoder-decoder models used in tasks like translation and summarization, you can start with the blog Transformer-Based Encoder-Decoder Models.

Model Architecture

In Figure 2-3, you see a high-level view of a typical decoder-only Transformer architecture. It shows how a prompt like “Write a short introduction about the US capital city” is processed by the model to generate a token such as Washington. The model can be conceptually broken down into three key components: tokenizer, transformer (decoder) blocks, and language modeling (LM) head. Let’s examine these components one by one.

Figure 2-3. Decoder-only transformer model architecture. The model is made up of a tokenizer and embedding, a list of transformer blocks, and an LM head.

Step 1: Tokenizer and Embedding

The first step in the LLM token generation is handled by the tokenizer, which takes raw input text (prompt) and does the following:

Breaks it into discrete tokens based on a fixed vocabulary.

Converts those tokens into token IDs, which are numerical representations. These IDs map to the actual text in the vocabulary used in step 1.

Maps those token IDs into dense vector embeddings using an embedding layer, preparing the input for the Transformer blocks.

This tokenization and embedding step ensures the model can interpret and process human-readable input in a form that the following transformer blocks will understand.

To better understand the tokenization process, you can explore it using OpenAI’s tokenizer tool. For example, given the input text “Write a short introduction about the US capital city”, the tool breaks it into 11 tokens: ["\"", "Write", "a", "short", "introduction", "about", "US", "capital", "city", "\""] . It then maps each token to a corresponding token ID: [1100, 10930, 261, 4022, 22575, 1078, 290, 2952, 9029, 5030, 693]. Ultimately, each token ID is mapped to an embedding vector (a numerical array). For example, token ID 1100 might correspond to a vector like [-0.12, 0.55, 0.98, ...].

Step 2: Transformer (Decoder) Blocks

The heart of the model lies in a stack of decoder blocks. These blocks are where most of the computation (and learning) happens during model serving and training.

The decoder blocks are stacked (for example, 12, 24, or more layers), allowing the model to build up a rich contextual understanding of the prompt and previous outputs—token by token—before predicting the next one.

The outputs of these decoder blockers are hidden states–contextualized representations of the input tokens. More specifically, they are tensors of shape [N, d], where N equals the number of tokens in the input sequence and d equals the hidden state size (for example, 768, 2048, 4096, depending on the model).

These hidden states (tensors) are then passed to the next component—the LM head. In most cases, only the final hidden state, corresponding to the last token in the sequence, is used to predict the next token in the output.

Step 3: Language Modeling (LM) Head

When the LM Head generates the output token, it primarily does two things:

First, it maps the hidden states from transformer blocks to vocabulary logits–a probability distribution for the tokens in the token vocabulary (see the Tokenizer section).

Second, it selects the output token based on the vocabulary logits (probabilities). Usually the token with the highest probability is selected as the next output. For example, in Figure 2-3 the LM head selects Washington , followed by alternatives like London, New York, and Cat.

With this foundational understanding in place, let’s now examine the architecture of a real-world model—Qwen 2.5—with the following code as a concrete reference.

model_name = "Qwen/Qwen2.5-0.5B"
# load Qwen2.5
model = AutoModelForCausalLM.from_pretrained(
   model_name,
   trust_remote_code=True,
   device_map="auto"
)
# Print all model configuration parameters
config = model.config
print("\n=== Model Configuration Parameters ===")
# Architecture parameters
print("\nArchitecture Parameters:")
# Size of the hidden layers
print(f"Hidden size: {config.hidden_size}")
# Number of transformer blocks
print(f"Number of layers: {config.num_hidden_layers}")
# Number of attention heads
print(f"Number of attention heads: {config.num_attention_heads}")
# Size of the MLP intermediate layer
print(f"Intermediate size: {config.intermediate_size}")
# Tokenizer parameters
print("\nTokenizer Parameters:")
# Size of the vocabulary
print(f"Vocabulary size: {config.vocab_size}")
# Maximum sequence length
print(f"Maximum position embeddings:{config.max_position_embeddings}")
# Model-specific parameters
print("\nModel-specific Parameters:")
for key, value in config.to_dict().items():
   if key not in ['architectures', 'model_type', 'torch_dtype']:
       print(f"{key}: {value}")

Before working with a model—especially an LLM—it’s a common and recommended practice to inspect its configuration. A model’s configuration contains valuable details about its architecture, such as the number of layers, hidden dimensions, attention heads, vocabulary size, and more. Understanding this information helps you estimate the compute resources required (such asGPU memory), choose appropriate model-serving strategies (like quantization or batching), andplan for performance optimization (for instance, layer-wise parallelism or model sharding).

Here is Qwen 2.5’s model configuration:

>> Architecture Parameters:
Hidden size: 896
Number of layers: 24
Number of attention heads: 14
Intermediate size: 4864

>> Tokenizer Parameters:
Vocabulary size: 151936
Maximum position embeddings: 32768

>> Model Size:
Total parameters: 494,032,768

>> Model-specific Parameters:
vocab_size: 151936
max_position_embeddings: 32768
 .. .. ..

By inspecting the configuration up front, you can make better decisions about deployment strategies—like whether to use quantized models, distillation, multi-GPU setups, or even offloading techniques.

Transformer (Decoder) Block

Next, let’s zoom into the Transformer blocks. Decoder blocks are the most critical and computation-heavy layers in the LLM architecture.

Figure 2-4. A Transformer (decoder) block is composed of a self-attention layer and a feedforward layer.

In Figure 2-4, you can see that a Transformer block contains two primary components:

Self-attention layer

This is the novel component of Transformer models, which incorporate contextual information by dynamically correlating input tokens based on their meaning and relative positions. Calculating attention (for context information) is crucial in LLM, since context allows models to understand the meaning of words and sentences based on surrounding information, reducing ambiguity and enabling coherent interpretation.

In simple terms, the self-attention layer allows the model to consider all previous tokens in the prompt and assign different levels of importance to each when generating the next word. In the next section, we’ll take a closer look at how attention works.

Feedforward neural network (FFN)

FFN refines the understanding of a given token into a stronger representation–the contextual meaning of a token, usually a dense, high-dimensional vector (for example, 768 or 4096 dimensions). This component applies per-token transformation. It leverages the context information from the attention layer, incorporates all the knowledge it learned during training, and provides a refined token representation before making a prediction.

The following code shows how to inspect the decoder layer of the Qwen 2.5 model:

print("\nModel Structure:")
def print_module_structure(module, prefix=''):
   for name, child in module.named_children():
       # Skip certain internal modules for clarity
       if name in ['_orig_mod', 'wrapped_model']:
           continue
       # Print the current module
       print(f"{prefix}{name}: {type(child).__name__}")
       if "Qwen2Attention" in name.lower():
         print(f"\nFound attention module: {name}")
         print(f"Type: {type(module).__name__}")
         # Print attention-specific attributes
         if hasattr(module, 'num_heads'):
             print(f"Number of attention heads: {module.num_heads}")
         if hasattr(module, 'head_dim'):
             print(f"Head dimension: {module.head_dim}")
         if hasattr(module, 'hidden_size'):
             print(f"Hidden size: {module.hidden_size}")
         if hasattr(module, 'rotary_emb'):
             print(f"Has rotary embeddings: {module.rotary_emb is not None}")

With the code, you can clearly see the structure of the decoder layer, including attention:Qwen2Attention, mlp:Qwen2MLP, and layernorm:Qwen2RMSNorm:

Model Structure:
model: Qwen2Model
  embed_tokens: Embedding
  layers: ModuleList
    0: Qwen2DecoderLayer
      self_attn: Qwen2Attention
        q_proj: Linear
        k_proj: Linear
        v_proj: Linear
        o_proj: Linear
      mlp: Qwen2MLP
        gate_proj: Linear
        up_proj: Linear
        down_proj: Linear
        act_fn: SiLU
      input_layernorm: Qwen2RMSNorm
      post_attention_layernorm: Qwen2RMSNorm
    1: Qwen2DecoderLayer
      self_attn: Qwen2Attention
    .. .. ..

Knowing this decoder layer structure is crucial when planning model optimization or serving. For example, attention-heavy blocks may benefit from fused attention kernels or custom CUDA implementations, which we will discuss in the following chapters.

Capture Token Context by Calculating Attention

Context is critical for language processing tasks. For example, given the sentence “I saw a dog chasing a squirrel, and it climbed up the tree” without context, it’s unclear whether it refers to the dog or the squirrel.

Ever since the 2017 publication of “Attention Is All You Need,” self-attention has become the main attention mechanism that transformers allow each token in a sequence to “look at” other tokens and weigh their relevance when computing its own representation.

For example, when an LLM processes the prompt “Write a short introduction about the US capital city,” at the token capital, the self-attention allows the model to:

look back at US to understand that capital is referring to a country’s capital, not financial capital.

consider introduction as a direction to keep the writing concise and general.

use Write to know the overall task is instructional.

Attention Calculation

In a nutshell: for each token, the LLM computes three vectors: query (Q), key (K), and value (V). The attention score of a token is calculated by taking the dot product between the query of the token and the keys of all the tokens of the input sequence, scaling the result, applying a softmax to get weights, and then using those weights to compute a weighted sum of the value vectors. This result becomes the updated representation of the token, enriched with context from the entire sequence. The attention mechanism is summarized by the following equation:

Attention(Q,K,V)=softmax(QKTdk)V

For a detailed explanation of this calculation, please refer to section 3.2.1, “Scaled Dot-Product Attention,” in “Attention Is All You Need.”

Multi-Head Attention

To further enhance the model’s understanding, instead of doing one attention calculation per token, the model performs multiple attention calculations for the same token, called heads, each with its own Q/K/V projections. This allows the model to capture different types of relationships (for example, syntactic, positional, and semantic) in parallel. The outputs from all heads are then concatenated and passed through a linear layer to produce the final attention output. Figure 2-5 illustrates a highly abstracted view of the multi-head self-attention mechanism described here.

Figure 2-5. Conceptual diagram of a multi-head self-attention calculation

Now let’s see how attention works in practice. In the following code, for the input sequence “write a short introduction about the US capital city”, we use the bertviz library’s “head view” class to visualize the attention of each decoder layer of the Qwen model. You can try this code yourself to look at how different attention heads in a layer connect different tokens:

text = "write a short introduction about the US capital city"
model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name, output_attentions=True)
model = AutoModelForCausalLM.from_pretrained(model_name, output_attentions=True).eval().cuda()
# Tokenize input
inputs = tokenizer(text, return_tensors="pt").to(model.device)
# Generate outputs with attention
with torch.no_grad():
   outputs = model(**inputs, output_attentions=True)
# Get attention weights
attention = outputs.attentions
# Use bertviz and attention weights to visualize
# the relations between the input tokens
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
head_view(attention, tokens)

By specifying output_attentions=True, we could fetch the attention weight of the given model input, “write a short introduction about the US capital city” after the token generation. And then we could visualize tokens relations in each decoder layer by head_view(attention, tokens). See the output graph in Figure 2-6.

Each line in Figure 2-6 represents an attention connection between two tokens—the line thickness indicates the strength of the attention weight, while color distinguishes different attention heads. By examining these connections, we can observe which tokens a given token attends to most strongly. In this example, the token capital shows strong attention toward US and write, suggesting these tokens are influential in shaping its contextual representation at layer 10.

Figure 2-6. Multi-head attention visualization for “capital” Token in Transformer Layer 10

It’s important to note that the attention computation described here is intentionally high-level and conceptual. We’ve abstracted away many implementation details—such as token and embedding dimensions, matrix shapes within each attention head, and the exact mechanics of output aggregation. This simplification is deliberate: for ML engineers focused on model serving and optimization, the current level of understanding—centered on how attention works and how tokens are generated in a Transformer—is typically both sufficient and practical.

However, if you are interested in exploring the attention mechanism in greater depth, I highly recommend The Illustrated Transformer blog, which walks through attention calculations in an intuitive, non-mathematical way, as well as books such as Hands-On Large Language Models and Build a Large Language Model (From Scratch), which offer detailed and comprehensive explanations of transformers and attention.

You just need to grasp the Transformer concept–not the math

Many people find the mathematical details of attention and Transformers intimidating—especially the matrix-heavy computations often emphasized in academic papers. However, for ML engineers focused on LLM serving, it’s both practical and—often—preferable to understand the attention mechanism conceptually, without diving deep into the math.

While a detailed understanding of self-attention can be valuable, it is not a prerequisite for effectively serving or optimizing LLMs. From a serving perspective, what’s most important to know is that attention is compute-intensive, especially during the prefill phase (which we will discuss shortly), and that its memory and latency costs scale with the input sequence length. This level of understanding is critical when you’re tuning inference performance, implementing KV caching, or distributing workloads across GPUs.

Now that you’ve built a solid understanding of LLM architecture and the autoregressive generation process, let’s shift to a more practical perspective. In the next section, we’ll walk through how to execute LLM generation step by step, using hands-on code examples to illustrate each stage of the process.

Executing LLM Generation: Step-by-Step Walkthrough

In this section, we’ll take you behind the scenes of how a LLM generates tokens—step by step. We believe the most effective way to understand the performance challenges of LLM serving is through hands-on, incremental execution, observing how tokens are generated one at a time.

By the end of this walkthrough, you’ll have a clear and practical understanding of how LLM inference works. More importantly, you’ll be introduced to three core concepts critical to high-performance LLM serving: the KV cache, prefill, and decoding. These concepts form the foundation for optimizing latency, throughput, and memory efficiency in modern LLM systems.

Run the Qwen model

To get started easily, let’s run the Qwen 2.5 model using the Hugging Face pipeline library. This is the simplest way to use an LLM, as the pipeline object abstracts away much of the underlying complexity and provides a straightforward generator() API for text generation.

# Initialize the text generation pipeline
generator = pipeline('text-generation', model='Qwen/Qwen2.5-0.5B')
# Define your prompt
prompt = "Write a short introduction about the US capital city."
# Generate text
generated_text = generator(prompt, max_length=50, num_return_sequences=1)
# Print the generated text
print(generated_text[0]['generated_text'])

This code loads the Qwen 2.5 model and generates text using generator(prompt, ...). For the input prompt “Write a short introduction about the US capital city,” the output is:

Write a short introduction about the US capital city. The United States of America is the largest country in the world by area, .. .. ..The city is also home to many other important institutions, including the National Mall, the White House Museum, and the National Archives.

Model prediction, line by line

Now, let’s take a look at what happens under the hood of the generator() function. Before diving into the code, let’s first revisit the token-by-token generation workflow, illustrated in Figure 2-7.

Figure 2-7. The token-by-token generation workflow in LLMs.

Figure 2-7 highlights two key aspects of LLM generation. First, the LLM generates one token at a time; second, each newly generated token is appended to the previous input sequence to form the new input for the next prediction step.

Now, let’s examine how this process is implemented. First, we need to specify the model and load tokenizer. In Example 2-1, instead of using pipeline, we use AutoModelForCausalLM to load the Qwen 2.5 model. It offers full control over the generation configuration, inputs, and decoding strategies—making it the ideal choice when we want to manually prepare inputs and handle the generation process step by step.

Example 2-1. Implementing Qwen token generation, token by token

# (1) load tokenizer and model
model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to("cuda" if torch.cuda.is_available() else "cpu")

Next, define the input prompt:

# (2) Define the input prompt
prompt = """The history of human … … … How might the next wave of communication tools shape our relationships, societies, and sense of identity?"""

Now, it’s time to tokenize the prompt–that is, convert it to an input format that the model understands. idx = tokenizer(prompt,..) tokenizes the prompt as the initial model input.

# (3) Convert (Tokenize) prompt to the input format that model understands
max_new_tokens = 100
# tokenize the input prompt for the first output token
# PS: prompt is the initial input sequence for LLM generation
idx = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
start_time = total_time = time.time()
times = []

Then we start the main generation loop, generating tokens one by one. Notice that outputs = model(idx_cond), at step 4(B), runs the Qwen model and produces candidates for the new token, given the input idx_cond, while logits = outputs.logits gets the logits, or raw prediction scores, for each token prediction:

# (4) Main generation loop - generate tokens one by one
for _ in range(max_new_tokens):
   # (A) Set the current context for generation
   idx_cond = idx
   with torch.no_grad():
       # (B) Generate predictions (token candidates) for next token
       outputs = model(idx_cond)
       # Get the logits (raw prediction scores) for each token prediction
       logits = outputs.logits

Next, the model selects the next token from the predictions it just generated:

# (C) Select next token from the predictions generated in step (B)
   logits = logits[:, -1, :] #  Select only the logits for the last token
   # Convert logits to probabilities using softmax
   probas = torch.softmax(logits, dim=-1)
   # Sample the next token from the probability distribution of the
   predicted tokens from step (B)
   idx_next = torch.multinomial(probas, num_samples=1)
   print("Next Token is:", tokenizer.decode(idx_next[0])
   time_cost = time.time() - start_time
   times.append(time_cost)

Append the new token to the input sequence using idx = torch.cat((idx, idx_next), dim=1). Finally, check that an end-of-sequence token was generated; idx also holds the final output text.

  # (D) Append the new token to the input sequence
   idx = torch.cat((idx, idx_next), dim=1)
   # (E) Check if end-of-sequence token was generated
   if idx_next.item() == tokenizer.eos_token_id:
       print("\n[Generation completed - EOS token reached]")
       break

Finish by decoding the entire generated sequence:

# (5) Decode the entire generated sequence
generated_text = tokenizer.decode(idx[0], skip_special_tokens=True)

Note

You can find the full sample code at ch2_Workthrough_LLM_execution.ipynb

You can try out this code yourself in the ch2_Workthrough_LLM_execution.ipynb notebook. In addition to demonstrating the code, we also tracked the latency of each token during the generation process. The demo generated 100 tokens in 9.1205 seconds—averaging approximately 0.09 seconds per token. See Figure 2-8 for details.

Figure 2-8. Time distribution for each token in the text generation.

You may also notice a pattern in Figure 2-8: except for the first token, the per-token generation time tends to gradually increase. This is expected—since each newly generated token is appended to the input sequence, the model must process a longer context with every step. As a result, both computation time and resource usage grow with each token. See Figure 2-9 for a visual reference.

Figure 2-9. The LLM reprocesses the entire input sequence each time to generate a new token, leading to increased computation as the sequence grows.

As machine learning engineers, we can quickly recognize the inefficiency shown in Figure 2-8: the model recomputes attention for all previous tokens every time a new token is generated. This repetitive computation over known tokens is costly and unnecessary. Can we avoid recalculating the same attention for previous tokens and instead focus computation on the newly generated one?

This brings us to the next topic—the concept of KV Cache.

Enable the KV Cache to Boost Performance

Now, imagine if you could cache the intermediate results from all previous tokens—specifically, the key and value vectors computed by the attention mechanism. Then, during each new token generation step, you would only need to process the newly added token, rather than recomputing attention for the entire sequence. That’s what enabling the key-value (KV) cache does: this optimization stores the attention keys and values computed at each layer for previously generated tokens, allowing the model to skip redundant computations during decoding (see Figure 2-10).

Figure 2-10. Accelerating LLM generation with KV caching. Only the new token is processed while previously computed key-value pairs are reused, reducing redundant computation.

As shown in Figure 2-10, the KV cache shifts the LLM calculation from full-sequence recomputation to an incremental, cache-augmented workflow—trading increased memory usage for significant compute savings.This technique dramatically reduces redundant computation and enables much faster token generation.

With this understanding in mind, let’s look at how KV caching can be used with the code in Example 2-2.

Example 2-2. Implementing Qwen text generation, token by token, with KV caching enabled

# (1) Define Key/Value Cache for faster generation
past_key_values = None
for _ in range(num_interations):
   print("input_ids size: " + str(input_ids.size()))
   with torch.no_grad():
       outputs = model(input_ids=input_ids,
         # (2) Use KV-cache from previous iteration
         past_key_values=past_key_values,
         # (2) Enable KV caching
         use_cache=True,
         max_new_tokens = 100,
         min_new_tokens= 100)

       logits = outputs.logits
       # (3) Update KV Cache
       past_key_values = outputs.past_key_values
       torch.cuda.synchronize()

   logits = logits[:, -1, :]
   probas = torch.softmax(logits, dim=-1)
   generated_token_id = torch.multinomial(probas, num_samples=1)

   # (4) Update input_ids with only the new token (using KV-cache)
   # Note: Not concatenating with previous tokens due to KV-cache
   input_ids = generated_token_id
   idx = torch.cat((idx, generated_token_id), dim=1)

   if generated_token_id.item() == tokenizer.eos_token_id:
       print("\n[Generation completed - EOS token reached]")
       break

The biggest differences between Examples 2-1 and 2-2 are:

Example 2-2 only uses the newly generated token, generated_token_id, as the model generation input for the next token: input_ids = generated_token_id and outputs = model(input_ids=input_ids, .. ).

In Example 2-2, besides passing in the token, we also specify KV cache as the model generation input, see outputs = model(input_ids=input_ids,past_key_values=past_key_values, ..) at step 2.

Step 3 of Example 2-2 updates the KV cache after each token generation with past_key_values = outputs.past_key_values.

Our total execution time for Example 2-2 was 3.1416 seconds. Figure 2-11 compares the execution times over 100 tokens of Example 2-1 (left, no cache) and Example 2-2 (with cache).

Figure 2-11. Token generation time comparison: without vs. with KV caching enabled

Without caching, the generation time increases steadily as the model reprocesses the growing input sequence at each step. In contrast, enabling the KV cache significantly reduces compute time by reusing previously computed key-value pairs—resulting in a more stable and efficient generation speed after the first token.

With this example, you should now have a clear, hands-on understanding of how effective the KV cache is in improving LLM generation performance—both in terms of latency and throughput.

In practice, many advanced techniques have been proposed to design, manage, and optimize KV caches more efficiently. With the foundational intuition you’ve gained here, you’ll be better equipped to explore those methods and apply them in your own work. We’ll also cover KV cache optimization strategies in more detail in chapter 6 and 7, focused on performance tuning.

The Prefill and Decode Phases

Now that you’ve seen how tokens are generated, you’re ready to explore the next two phases, Prefill and Decode. These terms are widely used in the LLM serving domain, particularly in the context of performance optimization and scalability. Figure 2-12 illustrates the concept.

Figure 2-12. The Prefill and Decode phases in LLM text generation

Let’s look at these phases one by one:

Prefill phase

In the Prefill phase, also called prompt processing, the model processes the entire input prompt at once. This phase is compute-intensive because attention is computed across all tokens in the prompt (quadratic complexity with respect to sequence length).

Decoding phase

In the Decoding phase, also called token-by-token generation, the model generates one token at a time and repeats for every new token. The sequence grows, but the (attention) computation typically focuses only on the most recent token.

As Figure 2-12 shows, the Prefill phase occurs when the user submits the prompt “Write a short introduction about the US capital city.” During this phase, the model processes all tokens in the prompt simultaneously, leveraging parallelism for relatively efficient computation.

Once the prefill phase is complete, the Decoding phase begins. This is where the model generates tokens one at a time to form a response, such as Washington,D.C.,is, and the.

Figure 2-13 compares the difference in processing time between the Prefill and Decoding phases when KV caching is enabled.

Figure 2-13. Token generation time in the Prefill and Decoding phases of Example 2-2 (with KV cache enabled)

Figure 2-12 illustrates the performance difference between the Prefill and Decoding phases. The first bar on the left represents the time taken to generate the first token. This corresponds to the Prefill phase, which is significantly more time-consuming because the model must process the entire prompt at once. In contrast, the following bars (representing the Decoding phase) show much shorter and more consistent token-generation times, as the model generates one token at a time using the cached information from earlier steps.

Why does learning about the Prefill and Decoding phases matter?

Understanding the difference between Prefill and Decoding phases—and their respective compute patterns—is critical for making informed optimization and resource-planning decisions in LLM serving. The prefill phase is compute-intensive due to parallel processing multiple prompt tokens; the decoding phase is memory-intensive, primarily due to the frequent loading of model weights and the growing size of the KV cache. Knowing which phase dominates in your use case helps you target the right bottlenecks.

For example, in scenarios involving long prompts, such as processing a 500+ page PDF, the Prefill phase becomes expensive. On the other hand, if you’re working with short prompts and long generations—such as in chatbot replies or story generation—the Decoding phase becomes the bottleneck.

By distinguishing between these phases, you can better align your system design and optimization efforts with the performance characteristics of your specific application. We will discuss phase-specific optimization techniques in chapter 6 and 7.

In this section, we demonstrated how to implement inference from scratch—manually loading the model and generating tokens step by step. While this approach provides valuable insight into how LLMs work under the hood, real-world applications typically rely on serving frameworks rather than custom-built inference pipelines. In the remainder of this chapter, we’ll use the vLLM serving framework as a practical example to show you how to execute LLM inference with production-ready libraries.

Run the LLM with a serving framework

Model serving frameworks—such as vLLM and SGLang—are purpose-built systems designed to load pre-trained language models and expose them through APIs (typically REST or gRPC) for real-time or batch inference. While training frameworks like PyTorch and TensorFlow focus on model development and gradient-based learning, serving frameworks are optimized for efficient, scalable, and low-latency inference.

But serving frameworks offer far more than just the ability to run inference. They provide essential capabilities such as:

Efficient decoding with KV-cache reuse

Request scheduling (for example, batching or micro-batching)

Support for multi-user concurrency

Token streaming, cancellation, and interruption handling

A well-designed serving framework (like vLLM) also continually integrates the latest research optimizations—such as paged attention and speculative decoding—while abstracting away low-level infrastructure concerns. This allows developers to focus on building applications rather than re-implementing inference logic for each model architecture or chasing the latest academic papers about improvements.

In this section, we’ll walk you through how to use vLLM to load and serve an LLM model. You can compare this with the manual inference implementation we explored in the previous section to appreciate the practical advantages of using a serving framework.

Serve the LLM (Qwen) with vLLM

Serving an LLMwith vLLM is quite simple and efficient. With just a few lines of code, you can load a model such as Qwen2.5 and run inference. Here’s a basic example:

import time
from vllm import LLM, SamplingParams
model_name = "Qwen/Qwen2.5-0.5B"
# Load model with vLLM
llm = LLM(model=model_name, dtype="float16")

# Define the prompt.
prompt = """You are an expert AI historian writing a detailed chapter for a book titled "The Evolution of Human-AI Collaboration." … … Write in a formal tone, with rich detail and examples in each era."""

# Create inference parameters.
inference_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=128)

# Run token (text) generation with prompt and inference parameters.
outputs = llm.generate([prompt], inference_params)

# Print the results.
for output in outputs:
 print(f"Generated text: {output}")

As you can see, the LLM() and generate() functions abstract away much of the complexity, enabling quick experimentation with LLMs. Despite its simplicity, vLLM provides extensive configuration options to fine-tune performance and behavior to meet your specific application needs. Here’s an example with more advanced configuration:

# Configure Model Loading
model = LLM(
   model="Qwen/Qwen-7B",
   # Memory management
   swap_space=16,  # CPU swap space size in GB
   max_model_len=4096,  # Maximum model length
   # PagedAttention settings
   block_size=16,  # Block size for PagedAttention
   enable_prefix_caching=True,  # Enable prefix caching
   # KV cache management
   max_num_sequences=256,  # Maximum number of sequences
   max_sequence_length=4096,  # Maximum sequence length
   # Performance optimizations
   enable_chunked_prefill=True,  # Enable chunked prefill
   enable_cuda_graph=True,  # Enable CUDA graph
   # System settings
   worker_use_ray=False,  # Use Ray for distributed serving
   disable_custom_all_reduce=False,  # Disable custom all-reduce
)

# Configure Model Inference Parameters
sampling_params = SamplingParams(
   temperature=0.7,  # Controls randomness (0.0 = deterministic)
   top_p=0.9,  # Nucleus sampling parameter
   top_k=50,  # Top-k sampling parameter
   max_tokens=100,  # Maximum tokens to generate
   stop=["\n", "###"],  # Stop sequences
   frequency_penalty=0.1,  # Frequency penalty
   presence_penalty=0.1,  # Presence penalty
   repetition_penalty=1.1,  # Repetition penalty
   skip_special_tokens=True,  # Skip special tokens in output
)

vLLM follows the parameter conventions of the OpenAI text completion API, making it familiar and easy to adopt. You can find a full list of model-loading configuration options in the vLLM GitHub repository.

Performance Comparison: vLLM vs. Hugging Face Transformers

In addition to convenience, vLLM delivers substantial performance improvements. In many benchmarks, it achieves 10 to 20 times higher throughput than the Hugging Face .generate() API (transformer library). Let’s run a quick side-by-side comparison using the same prompt and model, see a Hugging Face Inference Example:

# Run model inference with Hugging Face library
# Load model and tokenizer in Hugging Face library
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B", device_map="auto", trust_remote_code=True)

start_time_basic = time.time()
# Create the model prediction pipeline.
generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
# Generate prediction with prompt
outputs_basic = generator(prompt, max_length=128, temperature=0.8, top_p=0.95)
end_time_basic = time.time()

vLLM takes 1.12 seconds, while the Hugging Face Library takes 19.58 seconds. That’s a 17x speedup on a single prompt—and the performance gap only widens when you scale up to concurrent or batched inference.

Best Practice: Start Simple, Then Optimize

In real-world development, teams can prototype with Hugging Face Transformers for ease of use, then migrate to frameworks like vLLM for production usage and fine tune the serving configuration for better latency, throughput, and concurrency. By allowing you to tune its serving parameters to fit your specific model and use case, frameworks like vLLM offer enterprise-grade performance with minimal engineering overhead.

LLM Streaming Serving Basics

In our vLLM code example, you might have noticed that although the LLM generates tokens one by one internally, the call to outputs = llm.generate([prompt], inference_params) waits until the entire output is generated before returning any result. In the context of a web service—for example, a chatbot—this means the user experiences a long delay before receiving any response–from a few seconds to even minutes, depending on the output length.

For applications with conversational interfaces (like chatbots) or those requiring real-time feedback (such as live captioning or summarization), this delay significantly impacts user engagement. A slow response can frustrate users, reduce retention, and ultimately, lead to a failure to meet business goals.

As you learned in the previous section, tokens are generated sequentially during the Decoding phase. To improve responsiveness, you can return each token immediately as it is generated—a technique known as LLM streaming. Streaming, in LLMs, refers to the process of incrementally returning output token-by-token (or in small chunks) during generation, rather than waiting for the entire output to complete.

Let’s look at how to enable streaming generation with vLLM using a simple code example.

# Initialize the vLLM async streaming engine arguments
engine_args = AsyncEngineArgs(
    model="Qwen/Qwen2.5-0.5B",
    dtype="float16",
    .. .. ..
)
# Create the vLLM async streaming engine
engine = AsyncLLMEngine.from_engine_args(engine_args)
# Define the async function to generate text in streaming mode
async def generate_text(prompt: str, max_tokens: int = 100):
    try:
        # Define sampling parameters
        sampling_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            stop=["\n"],  # Stop at newline
        )
        # Generate text
        request_id = "test-request"  # Unique ID for this request
        # Generate tokens in streaming mode
        results_generator = engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id
        )
        # Process the results
        final_output = None
        async for request_output in results_generator:
            final_output = request_output
            # Print each token as it's generated
            for chunk in request_output.outputs:
                print(chunk.text, end="", flush=True)
                # We could yield return the token chunk here
                # if it’s a web service
        return final_output

The key difference in the streaming version of the vLLM code is that we initialize the model using the AsyncLLMEngine class instead of the standard LLM class. With AsyncLLMEngine, the generate() function returns an asynchronous stream object (AsyncStream) that allows us to pull newly generated tokens one by one using an async for loop—for example, async for request_output in results_generator.

Another valuable benefit of streaming is that it allows users to cancel generation midway if the output is heading in an undesired direction. In vLLM, this can be achieved by calling engine.abort(request_id), where request_id is a unique identifier associated with the specific generation request. This feature not only improves the user experience by avoiding irrelevant or incorrect completions, it also helps conserve compute resources, which is especially important in production environments where efficiency and cost control are critical.

If you’re interested in exploring LLM streaming further, we’ve included a hands-on example in streaming.py. In the next chapter, we’ll demonstrate how to implement LLM streaming in a web service setting, enabling real-time interaction in user-facing applications.

LLM Batch Serving Basics

In the examples we’ve covered so far in this chapter, we’ve sent one prompt to the LLM at a time. While this approach is simple and effective for prototyping, it becomes a bottleneck in high-throughput scenarios—such as summarizing 100,000 documents, indexing 5,000 PDF files, or serving 20,000 of concurrent chatbot users. Processing one input at a time simply doesn’t scale.

To address this performance limitation and fully utilize GPU compute resources, you need batching. Batching refers to grouping multiple input requests together and processing them simultaneously in a single forward pass through the model. Instead of handling one prompt at a time, batching allows the model to generate outputs for multiple prompts in parallel, significantly improving throughput. While batching can be applied to most ML models, it is especially crucial for LLMs due to their compute characteristics. Figure 2-14 provides a conceptual illustration of how batching works in LLMs.

Figure 2-14. Batch LLM inference: processing input sequences and generating output sequences in parallel.

Transformer-based LLMs can process input batches efficiently because their computations—like matrix multiplications and attention—can be parallelized across sequences. Thanks to shared model weights and the GPU’s parallel compute capability, batching allows handling multiple requests simultaneously with minimal overhead.

We’ll explore advanced batching techniques in depth in chapter 6 and 7. For now, let’s walk through a simple experiment to see how batch generation works in practice and compare its latency to single-prompt generation.

# Prompts for batch generation, 4 input sequences
prompts = [
    "What is the meaning of life?",
    "Write a short story about a robot learning to love.",
    "Explain quantum physics in simple terms.",
    "Translate 'Hello, world!' into Spanish."
]
sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=100
)
start_time = time.time()
# process four input sequences together in one batch
vllm_outputs = llm.generate(prompts, sampling_params)
end_time = time.time()
vllm_time = end_time - start_time
print(f"\nvLLM generation time for 4 prompts in a batch: {vllm_time:.4f} seconds")
# process prompt one by one
start_time = time.time()
for prompt in prompts:
   vllm_outputs = llm.generate([prompt], sampling_params)
end_time = time.time()
vllm_time = end_time - start_time
print(f"\nvLLM generation time for 4 prompts one by one: {vllm_time:.4f} seconds")

When we ran this code experiment, vLLM took 1.0626 seconds to process a batch of four prompts; it took 2.3865 seconds in total to process all four prompts in one by one fashion. This demonstrates a 2.2x improvement in throughput when using batching over single-prompt execution.

With further optimization techniques, you can achieve even higher throughput while simultaneously reducing latency. One such technique is continuous batching (which we’ll discuss in detail in chapter 6), where new requests are dynamically added as others complete. For example, a 2023 study by Anyscale showed that continuous batching can improve LLM inference throughput by up to 23 times , along with significantly reducing p50 latency.

Summary

In this chapter, we took a deep dive into the core concepts, architectural foundations, and practical techniques for serving LLMs. We started by demystifying the Transformer architecture—focusing on decoder-only models—and examined how attention mechanisms and autoregressive token generation work.

From there, we moved into hands-on walkthroughs that showed how to execute LLM inference. You learned concepts like KV caching and the Prefill and Decode phases as critical components in optimizing LLM serving performance.

We then transitioned from scratch implementations to using modern serving frameworks like vLLM, demonstrating how they simplify deployment while delivering massive performance. These tools abstract away low-level complexity and enable developers to build scalable, production-ready LLM services with minimal overhead.

Finally, we explored two powerful strategies for real-world LLM application: streaming and batch.

Together, the techniques introduced in this chapter form a foundation on which you can build fast, scalable, and cost-efficient LLM services. You are now equipped with the essential tools and mental models you’ll need to reason about LLM performance, troubleshoot bottlenecks, and design more effective serving pipelines.

In the next chapter, we’ll build on this foundation and show you how to wrap these serving capabilities into a real-world LLM serving web service, exploring design decisions, API structure, and operational best practices.

¹ For simplicity, we use the terms input text, input sequence, and prompt interchangeably throughout the book to mean input to an LLM. We also use the term token interchangeably with word. However, in practice, a token is defined by the tokenizer algorithm, which determines how text is broken down into smaller units. Depending on the tokenizer and language, a token may represent a whole word, a subword, or even punctuation. On average, one token corresponds to approximately 0.75 words in English. You can explore the OpenAI tokenizer viewer to understand how a piece of text might be tokenized by a language model, and the total count of tokens in that piece of text.

Chapter 3. Building Model-Serving Services from Scratch

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 3rd chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

In this chapter, our goal is to help you build a strong intuition for implementing a model-serving service—one that exposes a machine-learning model, particularly LLMs, to end users via a web-based API.

In Chapter 1, we introduced the general paradigms of model serving, covering architectural patterns and common trade-offs. In Chapter 2, we examined how LLMs perform inference and generate text. Now, we turn to the engineering perspective: how to organize code and infrastructure to construct a complete serving stack—from scratch—for both single-model and multi-model serving scenarios.

Model serving is a rapidly evolving field, with hundreds of solutions available across open-source and commercial ecosystems. Navigating this landscape—evaluating, adopting, and customizing the right solution—can be overwhelming. We believe the most effective way to understand complex systems is through small, well-designed examples.

To that end, we’ve created two simplified yet representative sample services: one for single-model serving and one for multi-model serving. While they are intentionally streamlined, these examples capture the core components and architectural decisions needed to tackle the most common challenges in modern model serving. They offer hands-on insight into real-world requirements, helping you build a strong foundation for understanding, extending, and adapting model-serving systems.

We’ll begin by building a single-model LLM serving service that supports batching and streaming. After that, we will explore a common single-model serving design pattern. hen we’ll move on to implement a multi-model serving service, followed by an in-depth comparison of two design variations: one optimized for cost, and another for latency and scalability.

By the end of this chapter, we expect you have a solid understanding of what happens under the hood in both single and multi-model serving systems. These hands-on experience and design deep dives will enable you to evaluate, adapt, and extend open-source or cloud-based serving solutions to fit your specific use case with confidence.

Note

You can find the full sample code at this chapter’s accompanying GitHub repository. We’ve selected and simplified key portions of the code for demonstration purposes. For complete implementation details and accurate context, please refer to the full repository. Please also refer to the README file for step-by-step instructions on running the demo services locally.

Build an Online LLM Serving Service from Scratch

In this hands-on section, you’ll learn how to build an online single-model serving service tailored for LLMs from scratch. We’ll kick things off by reviewing the design goals and overall architecture of the service, then gradually move into feature development. You’ll first implement basic functionality, like handling a single generation request, then extend the service to support batching and streaming. Along the way, you’ll get a clear picture of the core challenges in serving LLMs efficiently. To wrap up, we’ll show you how to use model-serving frameworks such as vLLM to address some of the limitations we encountered during our custom implementation.

Design Goals

In this exercise, you’ll build a model-serving service that loads a single LLM at startup and supports concurrent generation requests for both batch and streaming. While this example is intentionally simplified to only work for one LLM model that can run on CPU, it covers all the essential components needed to scale into a multi-node, generalized production-grade system with advanced model optimization techniques in the future.

Rather than building a complex, production-ready service, our approach is to implement the core components to help you understand the following critical aspects of LLM serving:

How web APIs are designed to handle generation requests—including batching and streaming

What a typical LLM request-processing workflow looks like

How the service is structured internally to support flexibility and scalability of different LLM models

How concurrent requests are grouped and processed in batches

How streaming generation works under the hood

How batching and streaming can coexist in the same system

Where the key performance bottlenecks tend to appear in LLM serving

Let’s dive into the service architecture to get a high-level overview of how this sample service is structured and how its core components fit together.

Service Architecture

This sample LLM serving service has six core components:

API server

Deals with HTTP API endpoints and request/response handling

LLM engine

High-level orchestrator, operates inference end-to-end

Workload manager

Handles request queuing and batch management

Model executor

Manages processes and coordinates model-worker execution

Model worker

Executes model inference execution in its own process

Model manager

Loads and caches the model

Figure 3-1 shows how these components come together to form the complete service.

Figure 3-1. Single-model serving system architecture

As shown in Figure 3-1, the LLM engine acts like an orchestra conductor—setting the stage and coordinating how everything runs. It initializes all core components, including loading the LLM model, and then orchestrates how generation requests are handled.

Since the service is designed to manage concurrent web requests, with support for both batching and streaming, a workload manager is introduced to track the status of each prompt included in the generation request and determine the next batch prompts for the model to process. This is a critical point where different batching strategies can be applied to optimize the LLM throughput.

The model executor and model worker are the core components responsible for running model inference. The model worker handles model loading, hosting, and execution—running in its own separate process. In contrast, the model executor is responsible for initializing the model workers, setting up worker groups (including inter-worker communication), and making cross-process calls to trigger inference and retrieve generation results from the workers.

While this multi-process setup may seem like overkill for a simple example, it reflects a common and practical pattern in real-world LLM serving systems that LLMs typically run on GPU devices. It’s a standard practice to isolate model execution in dedicated processes bound to specific GPU devices. Meanwhile, the main web service remains on the CPU, focusing on handling concurrent request management and orchestration. This is because GPUs are expensive, and you don’t want GPUs running idle, waiting for CPUs to perform other tasks such as input pre-processing, output post-processing, tokenization, and de-tokenization. In a multi-process setup, all compute-intensive tasks are moved to different processes, leaving dedicated CPU resources for CPU-GPU interaction. This helps to ensure optimal GPU utilization.

With the architecture in mind, let’s move on to implementing the service. We’ll start by handling a single generation request, then progressively add support for batching and streaming.

Implement Single Generation Request Handling

In this section, you’ll implement a basic LLM model serving, handling one request (one prompt) at a time. First, initialize the core components in the LLMEngine class:

class LLMEngine:
   def __init__(self):
       self.model_executor = ModelExecutor()
       self.workload_manager = WorkloadManager()
       self.max_tokens = 20

# Initialize and set up the model worker
self.model_executor.setup_worker("facebook/opt-125m")

Next, set up the ModelWorker and ModelExecutor classes. In this example, you’ll run a single ModelWorker in its own process. The model executor communicates with the worker using two event queues: task_queue and result_queue. It sends prompt requests to the worker via task_queue and receives generation results from result_queue.

class ModelExecutor:
   def __init__(self):
       self.task_queue = mp.Queue() # Model request queue
self.result_queue = mp.Queue() # Model result queue

   def setup_worker(self, model_name: str):
       self.worker_process = mp.Process( # Single process worker
           target=ModelWorker.run,
           args=(model_name, self.task_queue, self.result_queue)
       )
       self.worker_process.start() # Start the worker process

Now, let’s prepare ModelWorker, which loads the “facebook/opt-125m” model from HuggingFace with AutoModelForCausalLM class with the help from the model manager.

class ModelWorker:
   def __init__(self, model_name: str):
       self.device = "cuda" if torch.cuda.is_available() else "cpu"
       self.model, self.tokenizer = ModelManager().load_model(model_name)
class ModelManager:
   def load_model(self, model_name: str = "facebook/opt-125m"):
       model = AutoModelForCausalLM.from_pretrained(model_name)
       tokenizer = AutoTokenizer.from_pretrained(model_name)

       return model, tokenizer

The worker runs a while loop in its own process (see the following run function), which continuously monitors the task_queue. When a request arrives in the queue, it pulls the prompt from the task_queue, performs model inference, and places the generation result into the result_queue. See the implementation as follows:

class ModelWorker:
   @staticmethod
   def run(model_name: str, task_queue: mp.Queue, result_queue: mp.Queue):
       worker = ModelWorker(model_name)
       while True:
           logger.debug("Waiting for batch from queue...")
           request = task_queue.get()
           result_queue.put(('complete', worker.generate(request)))

Now that the stage is set, let’s connect all the components to implement the actual model-serving logic. Refer to Figure 3-2 for an overview of the code execution workflow.

Figure 3-2. Model-serving workflow for a single generation request

You’ll start by defining your web API – the basic_generate endpoint. This endpoint accepts a single prompt as the request payload and returns the generated text as a plain string:

@app.post("/basic_generate", response_model=GenerateResponse)
async def basic_generate(request: GenerateRequest, \
                         llm: LLMEngine = Depends(get_llm)):
   generated_text = llm.basic_generate(request.prompt)
   return GenerateResponse(generated_text=generated_text)

class GenerateRequest(BaseModel):
   prompt: str
class GenerateResponse(BaseModel):
   generated_text: str

Next, you’ll implement the orchestration logic in LLMEngine. First, the basic_generate function sends the prompt to the ModelExecutor. Then the ModelExecutor passes the prompt to the ModelWorker via the task_queue. When the ModelExecutor receives the generated result from the result_queue, it returns the result back to LLMEngine, which returns it to the web API:

class LLMEngine:
   def basic_generate(self, prompt: str) -> str:
       sequence = Sequence(str(uuid.uuid4()), prompt, None, None)
	# Execute the prompt generation
	results = self.model_executor.execute(sequence)
       return results[0]['generated_text']

class ModelExecutor:
   def execute_batch(self, prompt: str):
       # Send prompt to the ModelWorker’s task queue
	self.task_queue.put((prompts, False))
       # Collect generation results from ModelWorker
	results = self.result_queue.get()
       return results

The final piece is implementing the inference logic within ModelWorker. In its run function, the ModelWorker picks up requests from the task_queue, performs model generation, and sends the results back to the ModelExecutor via the result_queue.

class ModelWorker:
   def generate(self, prompt: str):
       # Run model inference, model is loaded during service initialization
       outputs = self.model.generate(...)
       generated_text = self.tokenizer.decode(outputs[0], …)
       return {
           'request_id': prompt_data.id,
           'generated_text': generated_text
       }
   @staticmethod
   def run(model_name: str, task_queue: mp.Queue, result_queue: mp.Queue):
       … …
       request = task_queue.get()  # Fetch generation(prompt) task
       # Return result to ModelExecutor
	result_queue.put(('complete', worker.generate(request)))

With that, you’ve completed the first version of your online model-serving service. If you compare it to the LLM model execution we explored in Chapter 2, you can see how the model-serving service emphasizes building efficient software to expose model inference through a web interface.

Now you can use a few test cases to verify that the basic_generate API is working as expected:

def test_generate(client):
   response = client.post(
       "/basic_generate",
       json={"prompt": "Hello, I am"}
   )

If you roll out this service as-is, the first challenge you’ll encounter is low throughput. With the current inference service, users can send only one prompt at a time, and the model worker processes a single prompt per inference call.

As you learned in Chapter 2’s section “LLM Batch Serving Basics,” the compute resources are underutilized, since LLMs are well-suited for processing multiple prompts together in a batch. In the next section, you’ll upgrade the service to support batching, which will significantly improve throughput.

Batching

Now let’s implement batching using a new generate API. This API accepts a list of prompts in a single prediction request and returns a list of generated texts in the same order (to their prompt) as the input array.

@app.post("/generate", response_model=BatchGenerateResponse)
async def generate(request: BatchGenerateRequest):
   generated_texts = llm.generate(request.prompts)
   return BatchGenerateResponse(generated_texts=generated_texts)

class BatchGenerateRequest(BaseModel):
   prompts: List[str]

class BatchGenerateResponse(BaseModel):
   generated_texts: List[str]

With the API in mind, let’s discuss the general idea of batching before you implement it. Imagine you receive two requests: request1 contains two prompts (promptA and promptB), while request2 includes three prompts (promptC, promptD, and promptE). To handle these two requests in your service, you need to solve two key challenges. First, you need to combine prompts from different requests into a few batches, so the LLM can execute prompts in big batches, not per request, maximizing resource utilization. Second, you need a way to accurately map the generated outputs back to their corresponding original requests, so the service can return the results to the correct users in the right order.

Figure 3-3. How batching is handled in service design

Figure 3-3 presents our updated design, extending the single-request service architecture from Figure 3-2. To support batching, in this version, we’ve introduced a tracking data structure called Sequence for each prompt and a new component called WorkloadManager, and we’ve given the LLM engine additional responsibilities.

Follow the sequence in Figure 3-3, using the numbers in the diagram, to see how the system processes the two incoming generation requests with their multiple prompts.

Request Intake: The API server receives the requests and extracts individual prompts from each payload.

Prompt Queuing: The API server passes the requests to the LLM engine, which forwards them to the workload manager.

Prompt Tracking and Batching: The workload manager assigns a unique ID to each prompt and wraps it into a Sequence object. The workload manager maintains in-memory tracking of all active sequences and determines which prompts should be grouped into the next batch for processing.

Batch Execution: The LLM engine retrieves the next prompt batch from the workload manager and sends it to the model executor and worker for inference.

Model Inference: The model worker processes the entire batch in one inference call and returns the generated texts, each paired with its corresponding prompt ID.

Response Mapping: The LLM engine uses the prompt IDs to map each generated text back to its originating web request. This allows it to correctly return the generated outputs—for example, mapping GenTextA and GenTextB to request1, and GenTextC, GenTextD, and GenTextE to request2.

With this workflow in mind, let’s dive into the core implementation—starting with how each prompt is tracked and how we decide which prompts go into a batch.

Example 3-1. Workload manager batch scheduling implementation

# Each prompt is represented as a Sequence object
class Sequence:
   def __init__(self, seq_id: str, prompt: str, client_stream, loop):
       self.id = seq_id
       self.prompt = prompt
       self.output = []
       .. .. ..

class WorkloadManager:
   self.batch_size = 4 # Set max 4 prompts in a batch

   # Create a Sequence object for each prompt and assign an ID.
   def add_request(self, prompt: str) -> str:
       request_id = str(uuid.uuid4())
       sequence = Sequence(request_id, prompt, None, None)
       self.incoming_queue.put(sequence)
       self.sequence_map[request_id] = sequence
       return request_id

     # Use first-in-first-out (FIFO) strategy
     # to generate next batch of prompts
     def get_next_batch(self) -> List[Sequence]:
       while len(self.active_sequences) < self.batch_size \
             and not self.incoming_queue.empty():
          sequence = self.incoming_queue.get()
          self.active_sequences.append(sequence)

       return self.active_sequences

We want to highlight two key aspects of the get_next_batch function. First, the Workload Manager uses a first-in, first-out (FIFO) strategy to determine which prompts are selected for the next batch. Second, we’ve capped the maximum batch size, allowing the LLM to process up to four prompts at a time.

Throughput Optimization in Batching

In practice, batch size and batching strategy have a significant impact on inference throughput. Choosing the right batch configuration requires careful tuning and is far more complex than what’s shown in this simplified example. These settings must be adjusted often, based on the specific LLM model, the prompt characteristics, the nature of your web traffic, and the hardware on which the model is running.

In this example, we use a simple FIFO batching strategy to build intuition, but for a deeper exploration of batching optimization techniques such as dynamic batching and continuous batching, see Chapters 6 and 7.

Now let’s look at the implementation for batch-response mapping in LLMEngine.

Example 3-2. LLM engine batch implementation

class LLMEngine:
   # process multiple prompts in a request
   def generate(self, prompts: List[str]) -> List[str]:

       # Register prompt to workload manager
       # Track the prompt ID, assigned by workload manager
 prompt_ids = []
       for prompt in prompts:
           prompt_id = self.workload_manager.add_request(prompt)
           prompt_ids.append(request_id)

       # Keep executing batches until all prompts in the
       # requests are completed.
 while not self._is_batch_finished(request_ids):
           sequences = self.workload_manager.get_next_batch()
           results = self.model_executor.execute_batch(sequences)
           self.workload_manager.update(results)

       # Get generated result by using prompt_id
       generated_texts = []
       for prompt_id in prompt_ids:
           generated_texts.append(\
               self.workload_manager.get_sequence(\
                    prompt_id).output[0])
             self.workload_manager.\
                    remove_finished_sequence(request_id)

       return generated_texts

As shown in the generate function of LLMEngine (Example 3-2), the generate logic first collects all prompt_ids associated with a generation request. It then uses these IDs to retrieve the corresponding generated text. The function completes only after results for all prompts are available. While batching may introduce latency for individual requests—since some may need to wait—it significantly improves overall service throughput by optimizing resource utilization across requests.

Now that you’ve implemented batching logic, use the following test code to verify your implementation:

def test_generate_batch(client):
   # Test with multiple prompts
   test_prompts = [
       "Hello, I am",
       "The weather is",
       "I want to",
       "The best way to",
       "The most efficient way to"
   ]

   response = client.post(
       "/generate",
       json={"prompts": test_prompts}
   )

Tracking Every Prompt’s Execution

Tracking each prompt individually is a common technique in LLM serving. It decouples the user’s web request from the actual model execution, allowing the system to reorganize prompts for more efficient processing in the backend. This abstraction provides the serving service with the flexibility to apply optimization techniques such as dynamic batching and prioritization (discussed in chapter 6).

In this example, each prompt in a user request is converted into a Sequence object (as shown in Example 3-1), which becomes the fundamental unit for managing the LLM execution workload.

Streaming with batching

In this section, we’ll add streaming support on top of the batching mechanism we just implemented. Currently, the batching code waits until all prompts are fully processed before returning results, which can introduce significant latency and lead to a poor user experience.

As we discussed in Chapter 2, LLMs generate inferences one token at a time. To reduce the latency of waiting for the entire batch to complete, we want to stream tokens back to users as soon as they are produced—while still batching requests internally at each generation step. This allows us to maintain high throughput while providing a more responsive real-time experience for users.

Table 3-1 outlines the new streaming and batching user experiences we plan to support.

Table 3-1. New streaming user experience

Time
Action
Backend Batch
Streamed Tokens

T0

User A’s prompt arrives

[Prompt1: “Hello, I am”

]

[Prompt1]

Prompt1: “a”

T1

User B’s prompt arrives

[Prompt2: “I want to"]

[Prompt1, Prompt 2]

Prompt1: “a”, “student”

Prompt2: “see”

T2

User C’s prompt arrives

[Prompt3: “I like to"]

[Prompt1, Prompt 2,

Prompt 3]

Prompt1: “a”, “student”, “[end]”

Prompt2: “see”, “a”

Prompt3: “eat”

T3

A finishes, D joins

[Prompt4: “The best way to"]

[Prompt 2,

Prompt 3,

Prompt 4]

Prompt2: “see”, “a”

Prompt3: “eat”

Prompt4: “success”

In Table 3-1, you can see how the service processes prompts over time—from T0 to T3. Prompts [1–4] are received from four different users [A–D] at different times [T0-T3]. The service batches these prompts in the backend for model inference. It also tracks the generated tokens for each prompt individually. As tokens are produced, they are streamed back to users in real time. Once a prompt completes, it is removed from the backend batch to make room for new requests. For example, Prompt1 is removed at T3, since it’s completed at T2.

With this streaming user experience in mind, check out Figure 3-4 to get an overview of how to implement streaming with batching in our sample service.

Figure 3-4. Workflow for serving streaming with batching

Let’s walk through Figure 3-4 step by step:

Request intake: The API server receives two generation requests, each containing a single prompt. Since this is an asynchronous interface, it returns a Server-Sent Events (SSE) stream, which allows the client to receive tokens incrementally as they are generated.

Prompt initialization: The LLM engine creates a Sequence object for each prompt. This object has an output buffer (to store generated tokens) and an event queue (to stream tokens back to the client).

Batching loop: The LLM engine runs a background thread that continuously orchestrates batch generation in a loop. Each loop iteration represents a single batch step for token generation.

Model Execution: A batch of prompts is sent to the model executor and worker, which runs inference and generates one single token per prompt in the batch.

Token Collection: The model worker returns a list of newly generated tokens, each associated with its prompt ID.

Sequence Update: The LLM engine appends each new token to its corresponding Sequence—updating both the output and the prompt, which will then be used for the next round of token generation.

Streaming Tokens: The LLM engine pushes the new tokens to each prompt’s event queue, which signals the API server that tokens are ready to be streamed back to the client via the SSE stream.

We introduced some changes in this implementation compared to the batching implementation in Figure 3-3. The main changes we introduced are:

The generation API is now asynchronous, allowing users to receive token updates as soon as they are generated.

ModelWorker generates one token per inference step, rather than producing the entire output in a single call.

WorkloadManager tracks partial outputs and updates each prompt with newly generated tokens in real time.

Each prompt now has its own event queue, enabling LLMEngine to pass tokens from ModelWorker to the async API layer efficiently.

LLMEngine includes a dedicated batch-processing thread that orchestrates token-level inference across prompts.

Streaming API Implementation

It’s time for the core implementation. In the following code snippets, we’ll focus on how generated tokens are pushed back to the client using SSE for streaming.

First, in a background thread (requests_processing_loop), the LLMEngine continuously pulls batches of prompts from the Workload Manager and sends them to the ModelExecutor and ModelWorker to generate the next tokens. Here is the first part of the background thread code:

class LLMEngine:
   # Streaming process loop, runs in the background thread.
   def requests_processing_loop(self):
     while True:
       # Get the next batch of prompts.
active_sequences = self.workload_manager \
 				.get_next_batch(is_streaming=True)
       prompts = [{'prompt': seq.prompt, 'request_id': seq.id} \
                      for seq in active_sequences]
       # Call model executor/worker to generate tokens as a batch.
       tokens = self.model_executor.execute_forward_batch(prompts)

Next, in requests_processing_loop, once the LLMEngine receives tokens from the ModelExecutor, it routes them to the corresponding request thread in the API Server using the event queue (client_stream) associated with each prompt. See the implementation in Example 3-3.

Example 3-3. LLMEngine batch-processing loop implementation

def requests_processing_loop(self):
        .. .. ..
    # Stream tokens back to respective clients
    for token in tokens:
       seq = self.workload_manager.get_sequence(token['request_id'])
       if result['is_finished'] or seq.token_count > self.max_tokens:
          asyncio.run_coroutine_threadsafe(
             seq.client_stream.put(None), # push finish token to request thread
             seq.loop
          )
          seq.finished = True
          self.workload_manager.remove_finished_sequence(token['request_id'])
       else:
          asyncio.run_coroutine_threadsafe(
          seq.client_stream.put(   # push token to request thread
             json.dumps({ \
                "token": token['token'], \
                "sequence_id": token['request_id']})),
             seq.loop)

          self.workload_manager.update_sequence_output( \
             token['request_id'], token['token'])

In each web request’s handling thread, event_generator (see the following code), the LLMEngine creates an event queue for the incoming prompt and registers it with the Workload Manager: workload_manager.add_streaming_request(prompt, queue, loop). The request can now wait for tokens on this event queue asynchronously:

async def event_generator(self, loop, prompt: str):
   # Create an event queue for this client's request
   queue = asyncio.Queue()
   # Register the prompt with an event queue request to workload manager
   seq_id = self.workload_manager.add_streaming_request(prompt, queue, loop)
   while True:
      # Get next token from queue
      data = await queue.get()
      if data is None:  # End of stream
         break
      yield f"data: {data}\n\n"

Meanwhile, the background thread—requests_processing_loop (see Example 3-3)—continuously processes prompts in batches and dispatches generated tokens back to the corresponding event queues. This code assigns an event queue to each prompt to enable communication between individual request threads and the centralized batch-processing thread.

Finally, the API server exposes the asynchronous streaming API generate_stream, which waits for tokens from the LLMEngine and streams them to the client using SSE via a StreamingResponse with the text/event-stream content type.

@app.post("/generate_stream")
async def generate_stream(request: GenerateRequest, llm: LLMEngine = Depends(get_llm)):
   async def event_generator():
       loop = asyncio.get_event_loop()
       async for token in llm.event_generator(loop, request.prompt):
           # token = 'data: {"token": " a", "sequence_id":
           #   "8310f5e1-6f6f-480e-b2f9-c8144a12cc17"}\n\n'
          yield token

   return StreamingResponse(
       event_generator(),
       media_type="text/event-stream"
   )

With the streaming implementation in place, we can now test it using the following code:

@pytest.mark.asyncio
async def test_generate_stream(async_client):
   prompt = "Hello, I am"
   # Create a streaming request
   async with async_client.stream("POST", "/generate_stream", \
         json={"prompt": prompt}) as response:
       assert response.status_code == 200
       # Process the response stream directly
       async for chunk in response.aiter_bytes():
          # Convert bytes to string and split by newlines
          lines = chunk.decode().split('\n')
          for line in lines:
             if line.startswith("data: "):
             data = json.loads(line[6:])
             token = data["token"]
       .. .. ..

If the data flow in this example is still unclear, Figure 3-5 provides a visual overview of how tokens are generated and returned to the client using streaming and batching.

Figure 3-5. Token streaming async workflow

Figure 3-5 illustrates how the LLM processes tokens in batches while still streaming them back to individual clients in real time. The figure is designed to be self-explanatory, so you can walk through the sequence it depicts. The core idea is to track each prompt’s execution using a dedicated event queue and dispatch newly generated tokens to their respective channels, ensuring timely and orderly delivery to each client.

You have now implemented all the key features of a single-model LLM serving service, including single-prompt inference, batch processing, and streaming inference. These examples should give you a sense of the substantial engineering effort involved in making LLM models accessible to end users and applications in a generic and scalable way.

In practice, building a full-fledged, production-grade model-serving service—with all the complexity of thread management, event queuing, inter-process communication, and handling execution details across various model architectures—can be quite daunting. This is where model-serving frameworks like vLLM come into play.

Serving Models with vLLM

Model-serving frameworks like vLLM are designed to abstract much of the complexity involved in deploying and operating ML models at scale. They manage critical tasks such as model loading and execution across a variety of architectures and frameworks, stay up to date with advances in model optimization, and address engineering challenges related to concurrency, threading, batching, streaming, scheduling, and more.

In the following example, we’ll demonstrate how to build a single-model serving service using vLLM. You can use the code in the GitHub repo to follow along. The service architecture is illustrated in Figure 3-6.

Figure 3-6. Single-model LLM serving with vLLM architecture

Compared to the streaming design shown in Figures 3-4 and 3-5, the architecture in Figure 3-6 is much more straightforward. All heavy lifting is delegated to vLLM, LLMEngine simply initializes vLLM with the facebook/opt-125m model and forwards generation requests directly to it. Look at the implementation in Example 3-4 to see how vLLM is integrated into LLMEngine for handling batch requests.

Example 3-4. LLM serving code with vLLM

class LLMEngine:
   def __init__(self):
       # Initialize vLLM with a model
       self.vllm_model = VLLM(model="facebook/opt-125m")

   def generate_vllm(self, prompts: List[str]) -> List[str]:
       # Configure sampling parameters for vLLM text generation
sampling_params = SamplingParams(
           temperature=0.7,
           top_p=0.95,
           max_tokens=self.max_tokens
       )

       # Generate text for all prompts
       outputs = self.vllm_model.generate(prompts, \
                                          sampling_params)
       # Extract generated text from outputs
       generated_texts = [output.outputs[0].text for \
                                         output in outputs]

       return generated_texts

In Example 3-4, you can see that using vLLM lets you enable the batch inference with just 10 lines of code. This is why such serving frameworks are so prevalent in model-serving implementation.

To optimize model serving frameworks, customize them

Although the vLLM framework abstracts away much of the complexity of model serving, you still need to customize it, unless you are OK with the default values. That’s why it’s essential to understand what happens under the hood in the serving system.

In general, knowing how model-serving systems work enables you to unlock the full potential of your framework and tailor it effectively to your system’s requirements, traffic patterns, and service-level objectives (SLOs/SLAs).

Also, many of the model optimization techniques introduced in Chapters 6 and 7 correspond directly to configuration options in vLLM. A solid grasp of model execution and serving fundamentals will give you the intuition you need to tune these parameters effectively. For example, when you understand how batching affects latency and throughput, you can fine-tune settings like max_batch_size and max_num_seqs to strike the right balance between real-time responsiveness and batch efficiency. Similarly, knowing how decoding leverages prompt reuse will help you optimize GPU memory allocation to support more concurrent users.

A serving framework can run as a standalone web server

Many model serving frameworks, such as vLLM and SGLang, can be deployed in two different ways, depending on your use case and integration needs. You can embed your framework as a library within your custom serving application, allowing for tight integration and greater control over execution flow, or you can run it as a standalone web server, exposing a REST or streaming API that can be called by external clients or other services.

In our sample code, we adopt the library-based approach for simplicity and easier demonstration. However, real-world deployments often run the serving framework as a standalone web server mode. We will explore that pattern in more detail in the next section.

A General Design for Single-Model LLM Serving

Now that you’ve implemented the sample service and looked under the hood, we hope you now have a solid foundational understanding of how a single LLM serving service operates. In Chapter 1, we introduced general design patterns for single-model services. However, given the importance of the concept of single-model serving—especially in the context of modern LLM serving—we’ll now look closely at a design specifically tailored for LLMs.

Requirements for Single Model Serving

The requirements any production-grade model serving system must meet generally include:

Low latency

The system should respond quickly, completing model inference and returning results with minimal delay to maintain user engagement.

High throughput

The system should be able to handle a high volume of concurrent requests. Maximizing throughput—measured in queries or tokens per second—is essential for serving at scale, reducing operational costs, and improving resource utilization.

Scalability

The system must scale horizontally to accommodate fluctuating traffic, seamlessly scaling out during peak demand and scaling in during quieter periods.

Reliability and availability

A production model-serving system must be highly available and fault-tolerant to ensure consistent and uninterrupted service.

Resource efficiency and cost management

Serving ML models, especially LLMs on GPUs, is resource-intensive. Efficient utilization of hardware (CPU, GPU, memory) is critical to keep the cost per query manageable.

Observability

Robust monitoring and logging are essential for tracking key performance indicators (KPIs) such as latency, throughput, and error rates, as well as model-specific metrics like prediction confidence. This visibility helps you identify bottlenecks and meet SLOs.

In addition to these general serving requirements, LLM serving introduces several unique challenges:

Large model size and memory footprint

LLMs often require tens to hundreds of gigabytes of memory, demanding careful management of GPU resources and memory allocation.

KV cache management

To support long context windows and efficient stateful decoding, serving systems must manage key-value (KV) caches effectively across multiple requests and sessions.

Streaming responses

For interactive applications, it’s important to stream tokens back to the client as they are generated.

Concurrency and batching with variable-length workloads

LLM requests often have varying input and output lengths. Serving systems must intelligently batch and schedule these heterogeneous workloads to maintain high throughput and GPU utilization.

LLM serving requirements are continuously evolving

Most traditional deep-learning model serving is relatively stable and often allows the model to be treated as a black box. with standard engineering solutions for latency and scalability. In contrast, LLM serving demands more dynamic and model-aware handling.

Requirements for LLM serving frequently change as model architectures and capabilities evolve, and are often model-specific. LLM input and output lengths also vary in different user scenarios. For example, KV cache behavior varies depending on the size of the user input and the model’s attention mechanism (such as multi-head versus multi-query attention).

Therefore, when designing an LLM serving system, it is crucial to isolate its evolving components from the stable infrastructure, ensuring flexibility for model-specific needs as well as robustness in the core serving logic.

General Design

In this section, we discuss a high-level system design that addresses both the general model-serving requirements and the LLM-specific challenges we have discussed. The design is intentionally kept conceptual and abstract, without relying on any specific tools, platforms, systems, or libraries. This allows it to be broadly applicable across a wide range of environments, whether deployed on private infrastructure or in the public cloud.

The core design idea is to group the aforementioned serving requirements into three distinct areas and address each separately:

Service infrastructure management

Focuses on scaling, availability, monitoring, and efficient resource allocation.

Business logic handling

Covers customer use-case integration, request batching, and streaming responses.

Model serving performance

Targets serving latency, throughput, and LLM-specific performance optimizations.

Figure 3-7 illustrates an overall service architecture based on this separation of concerns.

Figure 3-7. Single Model Serving General Design

First, in part A of Figure 3-7, we encapsulate the model serving logic into a replicable unit—such as a Docker container or a Kubernetes pod—and delegate infrastructure-level responsibilities to a distributed compute system. This system manages key tasks such as horizontally scaling service replicas up or down based on traffic demand, restarting unhealthy instances, dynamically allocating resources, and exposing logging and monitoring interfaces.

These infrastructure-management capabilities are standard features in most modern compute platforms, including public cloud providers like AWS and Azure, as well as open source solutions like Kubernetes. By abstracting infrastructure concerns from the serving logic, we decouple model serving from low-level details such as availability, fault tolerance, and resource provisioning.

In this setup, end users do not interact directly with individual service instances. Instead, requests are routed through a load balancer, which transparently distributes traffic across available instances. This approach hides the complexity of internal scaling and enables elasticity without exposing backend details to the user.

Next, in Part B, within each model serving instance is a serving frontend component that handles the web service interface and manages the core business logic of model serving. This layer is responsible for:

Authenticating and authorizing requests from customer applications

Integrating with external or internal systems, such as user customer data, model metadata, audit logs, and payment systems

Downloading and setting up models

Managing model configuration and runtime context

Preprocessing inference requests, including request validation, normalization, and batching logic

Traffic control, such as rate limiting and logging

This component acts as the middle layer between customer-facing interfaces and the backend model inference engine. This setup enables secure, flexible, and context-aware integration with the business environments, for example the customer-facing chatbot, while cleanly separating customer logic from the backend inference operations handled in Part C.

Part C shows the serving backend, the component where actual model inference takes place. This component typically runs as a separate process with access only to the serving frontend. It is entirely focused on delivering high-performance model execution.

The serving backend is designed to understand and support a wide variety of model architectures and optimization strategies. In practice, developers often rely on mature, production-grade model serving frameworks such as vLLM and Triton, which offer:

High throughput and low latency inference engines tailored for LLMs

Advanced optimization techniques, including quantization, KV cache management, and continuous batching of variable-length sequences

Efficient GPU utilization, maximizing hardware performance at scale

Ongoing evolution and strong community support, ensuring alignment with the latest model advancements

Isolating model execution in this backend component enables a modular, scalable architecture that cleanly separates business logic from inference performance concerns.

In the next chapter, we’ll explore some real-world examples that demonstrate how this single-model serving paradigm works in practice. The next part of this chapter focuses on a sample multi-model service.

Build a Multi-Model Serving Service from Scratch

In this hands-on section, you will learn the essential concepts for building a multi-model serving service. Unlike a single-model service that hosts only one model, a multi-model service can serve multiple models simultaneously. This approach is especially useful when you have many smaller models and want to avoid deploying a separate service for each one—significantly reducing serving costs and operational overhead.

We’ll start by implementing a serving service with custom model backends to illustrate the core principles of multi-model serving. Then, we’ll transition to using a production-grade backend—NVIDIA Triton Inference Server—to demonstrate how these concepts are applied in real-world systems.

By the end of this exercise, you’ll gain the insight needed to develop your own multi-model serving solution, whether building from scratch or leveraging open-source frameworks and cloud-native tools. Don’t forget to follow along with the code in the GitHub repo.

Design Goals

In this design, we’ll build a model serving service that hosts three models: two Transformer-based language models and one image-classification model. All three will run on CPU. The goal is to use just enough core components to help you grasp several critical aspects of multi-model serving, including:

Cross-framework support

How models built with different machine learning frameworks (such as PyTorch and ONNX) and architectures (such as NLP and Vision) can be hosted under a unified system.

Unified API interface

How a single service can expose a consistent web interface (for example, a REST API) for different types of models—regardless of whether the underlying task is text generation or image classification.

Resource management

How to manage and allocate limited compute resources (CPU, memory, and optionally GPU) when hosting multiple models on a single server instance, including strategies like lazy loading and least recently used (LRU)-based model eviction.

Let’s start with a high-level overview of how the system is structured and how its core components interact.

Service Architecture

This design consists of the following five core components:

API server

Exposes HTTP endpoints and handles request/response processing.

Model manager

Manages the model cache and coordinates model worker lifecycles.

Model store

Stores and retrieves model metadata.

Model engine

Responsible for creating model worker instances based on metadata provided by the model store.

Model worker

Loads the model into memory and handles inference execution.

Figure 3‑8 illustrates how these components work together to form the complete multi-model serving service.

Figure 3-8. Sample multi-model serving service architecture

Now let’s walk through the typical serving workflow illustrated in Figure 3‑8:

Client Request: The client sends a prediction request to the API server specifying a model ID and input payload. For example:

model_id = "550e8400-e29b-41d4-a716-446655440000"
input_data = "This movie was great! I really enjoyed it."

Model Lookup: The API server’s request handler forwards the request to the ModelManager, which checks whether the model is already loaded in the model cache.

Metadata fetch (if not cached): If the model is not found in cache, ModelManager queries the ModelStore for the model’s metadata, which defines how to load and host the model. A sample metadata entry might look like:

{
  "id": "550e8400-e29b-41d4-a716-446655440000",
  "name": "distilbert-base-uncased-finetuned-sst-2-english",
  "type": "text",
  "framework": "transformers",
  "version": "1.0.0",
  "description": "Sentiment analysis model"
}

Model Worker Creation: ModelManager passes the metadata to the ModelEngine, which creates a ModelWorker responsible for loading the model into memory and preparing it for inference.

Worker Registration: Once the ModelWorker is initialized, the ModelManager registers it in the model cache. If the cache is full, the ModelManager will evict the least recently used ModelWorker to make room for the new one.

Inference Execution: The API server retrieves the appropriate ModelWorker from the ModelManager and calls it to perform inference using the client’s input data.

Response: The API server returns the inference result to the client.

As you can see, the multi-model serving workflow is conceptually straightforward: receive a model request, load the requested model if needed, and execute inference. The core engineering effort lies in building an unified interface to host and operate models trained in different architectures and frameworks, while effectively managing the model cache to optimize compute resource usage.

Core Implementation

Let’s look at some key implementations, starting with prediction request handling from the API server.

class PredictionRequest(BaseModel):
   model_config = ConfigDict(protected_namespaces=())
   model_id: str
   input_data: Any

@app.post("/predict")
async def predict(request: PredictionRequest):
   # Get model worker
   worker = model_manager.get_model_worker(request.model_id)
   # Make prediction
   try:
       result = worker.predict(request.input_data)
       return result

In the predict API implementation, the ModelWorker is retrieved from the ModelManager and then used to execute the model inference. Let’s take a closer look at how the ModelManager handles model retrieval:

class ModelManager:
   def __init__(self, model_store: ModelStore, max_models: int = 2):
       self.model_store = model_store
       self.max_models = max_models
       self.model_cache = OrderedDict()  # model id -> worker
	self.model_engine = ModelEngine()

   def get_model_worker(self, model_id: str) -> Optional[ModelWorker]:
       # Check if model is in cache
       if model_id in self.model_cache:
           # Move to end (most recently used)
           self.model_cache.move_to_end(model_id)
           return self.model_engine.get_worker(model_id)

       # Get model metadata
       model_metadata = self.model_store.get_model(model_id)
       if not model_metadata:
           return None

       # Check if we need to remove least used model
       if len(self.model_cache) >= self.max_models:
           # Remove least recently used model
           id, model_worker = self.model_cache.popitem(last=False)
           self.model_engine.delete_worker(id)
       # Create and cache new model worker
       self.model_cache[model_id] = \
                 self.model_engine.create_worker(model_metadata)
       return self.model_cache[model_id]

In the get_model_worker function, you can see that the primary role of the ModelManager is to maintain a model cache—controlling which models are loaded into memory and how many can be kept active at once. In this example, we use a simple LRU strategy: when the number of cached models exceeds the max_model_count, the least recently used model is evicted to free up memory.

Next, let’s take a look at how the ModelWorker loads and initializes models:

class ModelEngine:

   def create_worker(self, model_metadata: ModelMetadata) -> ModelWorker:
       if model_metadata.id not in self.workers:
           if model_metadata.framework == "transformers":
               self.workers[model_metadata.id] = \
                         TransformerWorker(model_metadata)
           elif model_metadata.framework == "torchvision":
               self.workers[model_metadata.id] = \
                         TorchVisionWorker(model_metadata)
           elif .. ..
       return self.workers[model_metadata.id]

In the create_worker function, the ModelEngine selects the appropriate type of ModelWorker based on the model’s framework (for example, TorchVision). This example implements two types of workers—TransformerWorker and TorchVisionWorker—to demonstrate how different model backends are required to support a variety of model types. Let’s take a closer look at the implementation of the _load_model and predict functions of TransformerWorker in the following code:

class TransformerWorker(ModelWorker):
   def __init__(self, model_metadata):
       self.tokenizer: Optional[AutoTokenizer] = None
       super().__init__(model_metadata)

   def _load_model(self):
       if self.model is None:  # Only load if not already loaded
           self.model = AutoModelForSequenceClassification.\
                        from_pretrained(self.model_metadata.name)
           self.tokenizer = \
               AutoTokenizer.from_pretrained(self.model_metadata.name)

   def predict(self, input_data: Any) -> Dict[str, Any]:
       # Tokenize the input
	inputs = self.tokenizer(input_data, return_tensors="pt", \
                                padding=True, truncation=True)
       with torch.no_grad(): # Run model inference
           outputs = self.model(**inputs)
       predictions = torch.softmax(outputs.logits, dim=-1)
       return {"predictions": predictions.tolist()}

Now, let’s look at the model metadata logic. From the previous ModelManager and ModelEngine code, you can see that the system relies on each model’s metadata to determine the appropriate ModelWorker for handling an inference request. In our simplified example, the model metadata is defined as follows:

class ModelMetadata(BaseModel):
   id: str
   name: str
   type: str
   framework: str
   version: str
   description: str

class ModelStore:
   def __init__(self, config_path: str):
       self.models: Dict[str, ModelMetadata] = {}
       self._load_config(config_path)

   def _load_config(self, config_path: str):
       with open(config_path, 'r') as f:
           config = json.load(f)
           for model in config['models']:
               self.models[model['id']] = ModelMetadata(**model)

In real-world applications, model metadata is typically stored in a remote database or metadata service designed specifically for managing models. However, here we use ModelStore class to host metadata, which we load from a local JSON file for simplicity. Here’s a sample from that JSON file:

"models": [
       {
           "id": "550e8400-e29b-41d4-a716-446655440000",
           "name": "distilbert-base-uncased-finetuned-sst-2-english",
           "type": "text",
           "framework": "transformers",
           "version": "1.0.0",
           "description": "Sentiment analysis model"
       },
       {
           "id": "7c9e6679-7425-40de-944b-e07fc1f90ae7",
           "name": "pytorch/vision:mobilenet_v2",
           "type": "image",
           "framework": "torchvision",
           "version": "1.0.0",
           "description": "Image classification model"
       },
       .. .. ..

Let’s recap how our sample service addresses the three key design requirements:

To support various types of models, we implement different ModelWorker classes, each designed to handle the loading and execution logic for a specific model type.

To accommodate different model input formats, the predict interface uses a generic input and output structure, which is agnostic to the model used in the prediction request. We assume the client is responsible for preparing (preprocessing) the input and interpreting (post-processing) the output, as they are expected to understand the model they are invoking.

For resource management, models are loaded into memory on-demand upon receiving requests. We use an LRU cache to manage memory consumption—evicting the least-used models when memory usage exceeds a predefined threshold.

In practice, maintaining backend support for loading and executing a wide variety of models, managing model metadata and configurations, and coordinating thread safety and concurrency can be complex and error-prone. For these reasons, it’s often more efficient to delegate these responsibilities to a dedicated multi-model serving framework. This allows you to focus your own efforts on integrating that framework into your business applications.

Next, let’s explore how to leverage the NVIDIA Triton Inference Server to replace the ModelEngine and ModelWorker components in our current example.

Using NVIDIA Triton as Model Server

In the “Multi-Model Service” section of Chapter 1, we introduced the general concept of multi-model serving and highlighted NVIDIA Triton as one of the most popular and widely adopted solutions. It offers a standardized, high-performance way to serve models built with different formats—such as PyTorch, TensorFlow, ONNX, and TensorRT—through a consistent HTTP/gRPC API.

In this section, we’ll extend our multi-model serving example to demonstrate how to integrate Triton as the model-serving backend. But let’s first take a quick look at how Triton works in Figure 3-9.

Figure 3-9. How to work with Triton Server

As Figure 3‑9 shows, the Triton Server runs as a web service and exposes two key types of APIs: a Model Management API for loading, unloading, and configuring models, and aModel Inference API for sending prediction requests.

It takes just a few straightforward steps to run model inference with Triton:

Copy the model (in a Triton-supported format) into the designated model repository directory on the server, such as /models/densenet_onnx/.

Load the model using Triton’s Management API. For example:

            curl -X POST
            http://localhost:8000/v2/repository/models/densenet_onnx/load

This command instructs Triton to load the densenet_onnx model from the repository if it exists and is properly configured.

Send an inference request using the Inference API. For example:

curl -X POST http://localhost:8000/v2/models/densenet_onnx/infer

This triggers prediction on the loaded model using the input data provided in the request payload.

Now, let’s see how to integrate Triton Server into the multi-model service.

Figure 3-10. Integrating Triton as a serving backend to multi-model service

Compared to the previous multi-model design (shown in Figure 3‑8), the updated architecture in Figure 3‑10 introduces two new components: Triton worker and Triton server. The Triton worker acts as a wrapper that handles model loading and inference requests through the Triton Inference Server, which runs as a separate web service in its own process.

In this design, we delegate model hosting and execution responsibilities completely to Triton, while retaining model cache management, model file handling, and the external web interface within the multi-model service layer.

Although simplified, this design illustrates a common pattern in practice: using a wrapper service for handling the client web requests, integrating dependent systems, and managing resources (such as memory and CPU/GPU) and the model file lifecycle, while offloading the core inference workload to Triton.

Now let’s look at the code implementation for TritonWorker, starting with the model initialization logic. The TritonWorker will load the given model into Triton through the Triton management API, using the _load_model function as follows:

class TritonWorker(ModelWorker):
   def __init__(self, model_metadata):
      self.triton_url = "0.0.0.0:8009"  # Triton server URL
	 self.client = httpclient.InferenceServerClient(\
                                    url=self.triton_url)

   def _load_model(self):
       """Load model to Triton through Triton management API"""
  load_url = f"http://{self.triton_url} \
          /v2/repository/models/{self.model_metadata.name}/load"

       response = requests.post(load_url)

Then let’s look at the inference code. In the following predict function, the logic first converts the input data to the Triton format (which is declared in the model’s configuration), then sends a prediction request to Triton Server (self.client.infer()) to run model inference.

def predict(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
   """Make prediction through Triton inference API
   Args:
     input_data: Dictionary containing input tensors for
                 the model
     Each key should be an input name and value should be a
        numpy array Example: {"data_0": np.array(...)}

     Returns:
       Dictionary containing output tensors from the model
         Each key is an output name and value is a numpy array
       """
    # Create input tensors
    inputs = []
    for name, data in input_data.items():
       if not isinstance(data, np.ndarray):
       # Convert list or other array-like data to numpy array
   shape = data["shape"]
         content = data["data"]
         # Explicitly set dtype to float32
         array = np.array(content, \
             dtype=np.float32).reshape(shape)
       else:
         # Ensure existing numpy array is float32
 	   array = data.astype(np.float32)
         input_tensor = httpclient.InferInput(name, \
                                     array.shape, "FP32")
         input_tensor.set_data_from_numpy(array)
         inputs.append(input_tensor)

    # Hardcode output name for DenseNet model
    output_name = "fc6_1"

    # Make inference request
       response = self.client.infer(
       model_name=self.model_metadata.name,
       inputs=inputs,
       outputs=[httpclient.InferRequestedOutput(output_name)]
    )

    # Get predictions and convert numpy arrays
    # to lists for JSON serialization
   predictions = {
       output_name: response.as_numpy(output_name).tolist()
    }
   return predictions

Finally, let’s clean up resources on the Triton Server by sending a model unload request. Resource cleanup is a critical part of any multi-model serving system, as it allows you to reclaim limited compute and memory resources and make them available for serving other models.

def __del__(self):
   """Cleanup: unload model when worker is destroyed"""
   try:
     unload_url = f"http://{self.triton_url} \
         /v2/repository/models/{self.model_metadata.name}/unload"
     requests.post(unload_url)

Now that we’ve walked through the self-developed multi-model serving system, you can refer to our GitHub repository for detailed service setup instructions and local execution steps.

Tradeoffs in Multi-Model Serving Designs

This section dives deeper into two key approaches to multi-model service design—cost-optimized and latency-optimized.

Challenges

You’ve seen that multi-model services are designed to co-host multiple models within a single serving instance or container, allowing GPU, CPU, and memory resources to be shared efficiently. They support dynamic model loading and unloading based on real-time traffic, enabling significant cost savings and improved price-performance.

Multi-model serving is particularly useful when you’ve produced many models but don’t need to use them all simultaneously. For instance, let’s say you have scheduled processing jobs that only run for 3 or 4 hours a day. Instead of dedicating a separate container to each model, you can load the model on demand, run the job, and then unload it, freeing up resources to serve other models or tasks. As another example, perhaps you need to train 1,000 models for 1,000 customers. Since not all models will be used at the same time, you can limit the system to host at most 200 models concurrently, then load models on demand as requests arrive.

While the cost-saving benefits of multi-model serving are clear, its two biggest real-world challenges both lie in user experience:

Cold start latency

When a request arrives for a model that isn’t currently loaded, the serving instance must download the model file, load it into memory, and possibly evict another model if the cache is full. This process can introduce several seconds—or even tens of seconds—of delay, degrading the user experience. In high-traffic scenarios, where many models are cold, this can lead to request timeouts and cascading failures in downstream applications.

Hot model scaling

When a particular model suddenly receives high traffic, its latency increases as the load grows. Scaling that model becomes nontrivial—replicating the model across instances and updating the routing layer is complicated by the fact that each instance has an independent model cache. This creates engineering complexity and introduces nondeterministic behavior in performance.

There is no one-size-fits-all solution. As engineers, we constantly make tradeoffs. The two multi-model serving approaches we examine next are each optimized for a different priority: cost and latency.

A cost-optimized multi-model design

Let’s first look at a cost-efficient multi-model architecture (Figure 3-11) that derives from the paradigm in Chapter 1 (see Figures 1-8 and 1-10), incorporating additional details and refinements based on the concepts explored in this chapter.

Figure 3-11. Multi-model serving design optimized for cost efficiency

The design in Figure 3-11 optimizes for cost efficiency, with a focus on sharing serving resources across a variety of models, while attempting to mitigate cold start latency and hot model scaling issues.

In this architecture, all logic related to model loading, memory management, and hosting is encapsulated within each multi-model serving instance (Parts B and C of Figure 3-11). The serving frontend handles web APIs, caching, and model management, while the backend is responsible for inference execution and resource optimization. In practice, most developers rely on model server solutions such as Triton for the serving backend.

A key part of this design is shown in Part A: the model service API and routing logic. This component maintains mappings between models and serving instances, allowing it to:

Route requests to instances that already have the model loaded, minimizing cold starts.

Track the number of replicas per model and distribute models across instances, to scale hot models horizontally.

Distribute model traffic efficiently on the backend by applying bin-packing strategies to load models onto the minimum number of servers, optimizing resource usage.

However, this approach is not without challenges. The system is reactive—it adjusts to traffic patterns after they emerge. This means it’s always playing catch-up with demand, and prediction latency increases unavoidably when there are sudden traffic surges. Managing the routing logic, scaling decisions, and cache state across instances introduces complexity, making operations, debugging, and maintenance more difficult.

A latency-optimized multi-model design

Let’s look at a multi-model approach that favors latency. One common strategy in distributed systems engineering is to trade off capacity for improved performance. In this design, rather than sharing resources across multiple models, we provision a dedicated resource group for each model. This over resource provision approach enables better scalability and significantly reduces latency in the cold-start scenario, which works very well for high-demand models. See Figure 3-12 for the service architecture.

Figure 3-12. Multi-model serving design optimized for latency

The primary difference between the designs in Figures 3-11 and 3-12 is that we replace the multi-model serving instances with dedicated single-model serving instance groups—one instance group per model (Part A).

Another key change is that models are no longer loaded on demand. Instead, they are pre-provisioned by a model-provisioning service (Part C). Before sending prediction requests, the client must first call this provisioning service to create the instance group for the target model. Once provisioned, the service updates the model-to-instance-group mapping in the routing map. When a client sends a prediction request, the model service API consults the routing map and forwards the request to the appropriate pre-provisioned instance group for that model.

This design excels in latency and scalability, especially for hot models, since each model has its own dedicated, always-on resource pool. There’s no cold-start delay, and models can scale independently. Operators can also define separate resource policies for different models, providing greater flexibility in resource management and operations. The architecture is also simpler, making it easier to maintain and troubleshoot than the more dynamic, cache-heavy approach in the cost-optimized version.

The main trade-off is cost efficiency. In many cases, it’s difficult to predict which models will receive heavy traffic. As a result, you may end up overprovisioning resources for underutilized models, wasting compute capacity. Still, this approach is highly effective when you’re serving multiple models with steady and predictable demand.

Ultimately, we hope this comparison conveys that once you understand the fundamentals of model serving, you can tailor your architecture to meet your specific goals—whether that’s cost, performance, or operational simplicity.

Note

Multi-model serving applies to LLM, too

While LLM serving is typically implemented using a single-model approach—primarily due to its substantial compute and memory requirements—the multi-model serving paradigm remains highly relevant for many LLM use cases. In cases such as prefix caching and routing and multi-Low-Rank Adaptation (LoRA) support, multi-model serving provides significant advantages.

For prefix caching and routing, requests with shared prompt prefixes can be routed to specific model replicas that have the corresponding KV cache already populated, reducing redundant computation and improving both throughput and latency.

Multi-model serving enables dynamic management and loading of multiple LoRA adapters on top of a shared base model, allowing for scalable, memory-efficient personalization across tenants or use cases.

We will discuss prefix caching in Chapter 7 and multi-LoRA in Chapter 10.

Summary

In this chapter, we walked you through the code and designs to build two model-serving services from the ground up—introducing both single-model and multi-model serving architectures. Our goal was to help you develop practical intuition and hands-on skills for designing and implementing systems that expose ML models, especially LLMs, via web-based APIs.

We began by reviewing a simplified, self-developed single-model serving system that supports both batch processing and streaming. This walkthrough covered key components such as the API server, workload manager, model executor, and model worker—highlighting how they work together to manage concurrent requests, execute inference, and return generated outputs in real time. We also introduced advanced techniques like prompt tracking, token streaming, and concurrent batching.

To contrast and complement our custom implementation, we demonstrated how to leverage vLLM, a production-grade model serving framework, to abstract away much of the complexity while still allowing fine-grained performance tuning. This comparison underscored the value of understanding what happens under the hood, so you can tune framework settings effectively and make the right architectural decisions for your use case.

Next, we transitioned to multi-model serving, showing how a service can dynamically manage and serve multiple models in a shared infrastructure. You built a custom system that supports on-demand model loading, model metadata management, and memory-efficient model caching using an LRU strategy. This architecture supports cross-framework models like NLP and vision, all under a unified API interface.

To further reduce complexity and improve scalability, we demonstrated how to integrate NVIDIA Triton Inference Server as the backend for multi-model inference, offloading core responsibilities like model hosting, execution, and optimization.

Finally, we explored two key multimodel serving designs: one optimized for cost efficiency through shared resources and dynamic model loading, the other optimized for latency and scalability. These design patterns illustrate common tradeoffs in model serving, balancing cost, performance, scalability, and operational complexity.

We hope you now feel more confident about both building and evaluating model-serving solutions. With this foundational understanding, you’re now well-equipped to tailor serving architectures to meet your application’s unique requirements.

In the next chapter, we’ll apply what you’ve learned to some real-world case studies—examining how to set up model-serving infrastructure with both public cloud and open source offerings.

Chapter 4. Model Serving Best Practices

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 4th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

In Chapters 2 and 3, we explored model inference from implementation to system design, giving you a clear view of how LLMs execute and how to build a serving service from first principles.

This chapter steps up a level to cover the best practices of modern model serving. We begin with agentic applications: we unpack the critical components of a working knowledge agent and show how LLM serving energizes intelligent planning, tool use, and multistep reasoning.

We then present a layered reference architecture for an enterprise LLM serving platform. This framework highlights the essential building blocks and key considerations that separate a hobby stack from a production-grade platform.

Next, we examine build-versus-cloud choices as a spectrum, not a binary switch. You’ll see how to mix open-source stacks with cloud-managed offerings; why the “right” choice changes as traffic, cost, and business requirements evolve; and how knowing how to build equips you to choose how much to customize at any point. In practice, most teams live in the middle, layering custom development onto cloud services, with the degree of customization shifting over time.

We close with a look at core performance metrics for LLM serving, such as latency and throughput, which lay the foundation for the optimization chapters that follow.

By the end of this chapter, you will:

Recognize how LLM serving powers agentic workflows

Understand an enterprise-ready, layered serving architecture and its trade-offs

Navigate the build-versus-cloud spectrum and make situational choices that can evolve

Measure what matters so you can improve performance with intent

Armed with these practices—and the grounding from earlier chapters—you’re ready to dive into the next section on serving optimization, where we turn metrics into concrete latency, throughput, and cost wins.

Model Serving in an Agentic World

In this section, we use a sample agent to demonstrate how you can employ LLMs to create agentic user experiences, where software and services operate both interactively and autonomously.

Agents are both powerful and complex. A single agent can perform sophisticated tasks, but in practice, agents are often organized into an agent network or platform, where multiple agents build on top of one another, collaborate, or specialize in different functions. This layered architecture enables remarkable autonomy but also introduces significant design and operational challenges.

To help you navigate model serving in the context of agent development, we focus here on the fundamentals. Using a bare-bones agent as an example, we illustrate how model serving supports the core mechanics of agent autonomy. With this foundation in place, you will be better equipped to explore and understand more advanced and multiagent systems on your own.

Defining agents

In this chapter, we define an agent as an autonomous, LLM-powered system capable of:

understanding high-level goals

reasoning about how to achieve them

selecting and invoking external tools or data sources

adapting and iterating based on intermediate results

ultimately producing final outcomes with minimal or no human intervention

Types of agents include:

Research

Reads academic papers, extracts key findings, and synthesizes a literature review.

Coding

Generates, reviews, tests, and debugs code, while managing deployment from a code repository.

Business Operations

Automates data entry, generates reports, and distributes summaries to stakeholders.

The fundamental difference between an LLM-powered agent and a traditional assistant program (such as a rule-based chatbot with precoded workflows) lies in autonomy. Traditional systems depend on predefined instructions from developers or step-by-step directions from users. In contrast, LLM-powered agents can execute complex workflows independently, thanks to their ability to:

interpret natural-language instructions

reason and plan actions

select and use tools

adapt to feedback from the environment or humans

maintain awareness of context and user preferences

In this section, our focus is on building your intuition for how LLM model serving underpins agentic applications. For readers interested in the design and implementation of single- or multiagent systems, we recommend the following books: Building Applications with AI Agents by Michael Albada (O’Reilly, 2025) and AI Agents in Action by Micheal Lanham (Manning, 2025).

A sample knowledge agent

In this section, we use a real example to illustrate how model serving can power a knowledge agent designed to query and analyze information from PDF files. Although this sample agent is intentionally simple, this example highlights a common architectural pattern for building agents on top of model serving. For the sake of simplicity, we’ll name this knowledge agent “Knowledge Agent.”

To maximize portability, Knowledge Agent avoids local model hosting and databases. Instead, it relies on OpenAI’s APIs for both LLM inference and embeddings, while storing all information in memory. The source code and setup instructions are available in the “KnowledgeAgent” folder of this book’s GitHub repository. Before you read the rest of this chapter, please take some time to read these instructions and set this Knowledge Agent up on your own machine so that you can follow along with our examples. To gain a deeper understanding, we encourage you to run the sample agent yourself and experiment with different types of questions.

The sample Knowledge Agent supports document querying, intelligent planning for complex questions, summarization, and document analysis. Its knowledge base is built from the PDF files stored in the local knowledge_files folder. You can add your own PDFs to this folder, and the agent will automatically process them at startup:

~/llm-model-serving/ch04/KnowledgeAgent/ ls knowledge_files
  5-Level Paging and 5-Level EPT.pdf
  A Brief Introduction to the SAL.pdf
  A Brief Tutorial on Database Queries.pdf
  .. .. ..

Once the agent is running, you can ask it a wide range of questions. These may be direct information queries, such as:

“What is 5-level paging and how does it work?”

“What are Patricia tries and how are they optimized?”

They might also be advanced summarization and analysis tasks, such as:

“How do these technologies relate to modern computing systems?”

“Create a detailed comparison between database query optimization and data structure optimization.”

The following example demonstrates the agent running locally.

~/llm-model-serving/ch04/KnowledgeAgent/ python agent.py
Your question: What are the main types of database queries
   discussed in the tutorial?
Response:
This tutorial by Lutz Hamel explores the key tools used in modern
relational database systems: database queries, data mining, and OLAP
   .. ... ..

Additional examples can be found in the Example Queries section of the README.

The agent’s design

Let’s begin by examining the structure of the Knowledge Agent. Figure 4-1 provides a visual overview of its key components.

Figure 4-1. Knowledge Agent system overview.

In Figure 4-1, the Knowledge Agent is powered by two OpenAI-hosted models: text-embedding-3-small (for embeddings) and gpt-4.1-nano (as the LLM).

The text-embedding-3-small model converts text into numerical vector representations that capture semantic meaning. These embeddings enable tasks such as semantic search, retrieval, and content matching.

The gpt-4.1-nano model serves as the agent’s reasoning and language generation engine. It interprets instructions, plans subsequent actions, and produces natural-language responses.

Note

For more details on embeddings (converting text into numerical representations), see the OpenAI Embeddings User Guide.

On top of the models, the Knowledge Agent is built from the following core components:

Knowledge Agent (Orchestrator)

The central controller that coordinates all components.

Retrieval-Augmented Generation (RAG) System

Processes PDFs, generates embeddings, and performs vector search.

Planner

Leverages the LLM to create intelligent execution plans for user queries.

Actions (Executor)

Carries out specific tasks such as querying, summarization, and analysis.

The agent’s internal workflow

Now, let’s examine the internal logic of the Knowledge Agent. Figure 4-2 illustrates the workflow for handling a user query.

Knowledge Agent’s internal workflow for responding to a user query. Note that the LLM Serving illustrated in the graph can be identical or distinct LLM models, depending on the action scenario.

Let’s walk through the workflow shown in Figure 4-2 step by step:

A user sends a question: for example, “Create a detailed comparison between database query optimization and data structure optimization.”

The agent first calls its Planner component to design an execution plan.

The Planner calls the LLM, which generates a sequence of actions tailored to the query and the agent’s available capabilities. For example, for the input from step 1, LLM will create a three-steps sequential plan: query_rag_with_context, generate_analysis, and generate_summary:

{
'plan': ['query_rag_with_context', 'generate_analysis', 'generate_summary'],
'reasoning': 'First, retrieving relevant contextual information ensures a comprehensive understanding of both topics. Then, generating an analysis allows for a detailed comparison highlighting differences and similarities. Finally, summarizing consolidates the key points for clarity. This sequence ensures a thorough and well-structured comparison.',
'estimated_steps': 3
}

Next, the ActionExecutor executes the planned actions (from step 3) in sequence, it will first work on action query_rag_with_context. You can find the implementation of all three actions in our code repo,in the actions.py file.

The query_rag_with_context action first triggers the Retrieval-Augmented Generation (RAG) system to locate documents relevant to the query from the knowledge base, then use the LLM to summarize the documents. The RAG system functions as a semantic search engine over indexed documents. A sample relevant doc returned by the RAG system looks like following:

Document 1 (Source: A Brief Tutorial on Database Queries, Data Mining, and OLAP (hamel-197-manuscript-final).pdf):\nA Brief Tutorial on Database Queries, Data Mining, and OLAP  Lutz Hamel Department of Computer Science and Statistics University of Rhode  Island Tyler Hall Kingston ...

The ActionExecutor executes the second action, generate_analysis. This action sends the user query and the documents from step 5 as context to the LLM.

The LLM produces a deeper comparative analysis based on the query and its context.

Finally, the ActionExecutor works on the last action, generate_summary. It calls the LLM to summarize the analysis from the previous step.

The agent consolidates the outputs from all actions and returns the final result to the user.

By walking through this workflow, you can see how the Knowledge Agent decomposes a complex query into steps, retrieves the relevant information, and synthesizes an autonomous response with the help of an LLM.

Agent autonomy

The sample Knowledge Agent demonstrates that LLM-powered agents differ fundamentally from conventional applications, which typically work through procedural execution: they execute commands through fixed logic and predefined workflows. Instead, an agent can operate with goal-driven autonomy: it can interpret the user’s intent, choose an approach, integrate multiple tools, and deliver results, all with minimal user intervention.

Defining actions

To enable autonomy, an agent defines a set of reusable actions—discrete operations that it can execute in response to a plan. Typical actions include summarizing documents, answering queries, and performing deeper analysis. In our sample agent, each action is implemented as an LLM call with a specialized prompt template. For example, the create_analysis_prompt action generates a structured analysis based on the user’s query and supporting context:

def create_analysis_prompt(self, query: str, context: str) -> str:
    """Create a prompt for detailed analysis."""
    prompt = f"""
You are an expert analyst. Please provide a detailed
analysis of the following question based on the
provided context.

Question: {query}

Context:
{context}

Please provide:
1. A comprehensive analysis
2. Key insights and findings
3. Relevant examples from the context
4. Any limitations or gaps in the available information

Analysis:
"""
   return prompt

In practice, agent actions are not limited to LLM prompts. Depending on the design, they may involve:

Tool calls

Running a web search, querying an API, or executing a database command.

System operations

Managing files, triggering workflows, or interacting with external services.

Reasoning steps

Intermediate computations or planning subroutines that do not directly involve a user-facing response.

These diverse actions allow an agent to combine reasoning, knowledge retrieval, and external capabilities into a coherent workflow.

Model Context Protocol (MCP)

One popular approach to enabling tool use in agents is the Model Context Protocol (MCP). MCP standardizes how LLMs interact with external tools and data sources. Instead of hardcoding integrations, MCP defines a consistent interface for the agent to discover tools available in its environment, call them with structured inputs, and incorporate the results back into its reasoning process.

This protocol simplifies development and improves reliability by decoupling tool definitions from the core agent logic. MCP also helps address the brittleness of ad-hoc prompt engineering by giving LLMs a predictable way to invoke tools. You can check out the book AI Agents with MCP by Kyle Stratis (O’Reilly, 2026) for more detail.

For clarity and focus, our sample agent keeps its actions simple—sticking primarily to LLM-based calls with prompt templates. While real-world agents often combine LLM reasoning with tool invocation through protocols like MCP, this stripped-down design highlights the fundamentals of how model serving supports autonomy.

Planning with LLMs

Rather than relying on predefined workflows, the agent’s Planner uses the LLM to interpret a high-level instruction, break it down into subtasks, and determine the optimal execution path. The agent orchestrates these steps automatically, freeing users from manual micromanagement, such as parsing PDFs, locating relevant sections, and summarizing results.

A simplified version of the planning logic is shown below:

def create_plan(self, query: str) -> Dict[str, Any]:
  """Create an execution plan for the given query using OpenAI."""
  logger.info(f"Creating plan for query: {query}")

  # Create planning prompt
  planning_prompt = self.llm_manager.create_planning_prompt( \
     query, self.available_actions)

  # Get plan from OpenAI
  plan_response = self.llm_manager.generate_response( \
     planning_prompt, temperature=0.3)

  # Parse JSON response
  plan = self._parse_plan_response(plan_response)
  logger.info(f"Created plan: {plan}")
  return plan

Retrieval-Augmented Generation (RAG)

In the previous section, we introduced the RAG system (Figure 4-2), which serves as a search engine for the agent to locate text relevant to a user query within its knowledge base. Since RAG is central to how agents operate, we will explore it in more detail here.

Why RAG?

Why do LLM-powered agents need an auxiliary system like RAG? LLMs, while powerful, have well-known limitations:

Their knowledge is limited and static, since their training data is frozen at a specific cutoff date.

They may hallucinate, or generate content that is fluent but factually incorrect.

They often lack specialized or up-to-date information, and thus have domain-specific gaps in their knowledge.

RAG mitigates these issues by augmenting the LLM with external knowledge retrieved at query time. Instead of relying solely on the model’s internal training data, RAG enables responses enriched with domain-specific and continuously updated information. This improves accuracy, reliability, and adaptability—particularly for knowledge-intensive tasks. As a result, RAG-equipped agents are far more reliable and adaptable than agents that rely on LLM generation alone. Figure 4-3 illustrates the two main workflows of a basic RAG system.

Figure 4-2. A basic RAG system composed of index building and query (retrieval) workflows

Index-building workflow (offline process)

The index-building workflow (Figure 4-3, part A) is an offline data process, typically scheduled to run periodically. It begins by cleaning and parsing raw knowledge documents—often in formats such as HTML or PDF—into plain text. This text is then segmented into chunks, typically around 1,000 tokens each.

Chunking (splitting large content into smaller chunks) is more than just a preprocessing step for embeddings. It defines the granularity of retrieval (how much text the system can consider at once) and serves as the stage where offline batch inference occurs. Instead of embedding entire documents on demand, the system precomputes embeddings for each chunk in bulk, making query-time retrieval much faster and more efficient.

Once chunking is complete, each segment is encoded into a dense vector representation using an embedding model. These chunk embeddings are then indexed and stored in a vector database, where they become the foundation for efficient similarity search during the online query workflow.

Fine (Small) versus Coarse (Large) Chunking

Smaller chunks provide higher precision, since retrieval focuses on narrowly relevant passages. However, they can lose context if split too aggressively. Larger chunks provide richer context but may dilute precision by including unrelated content.

In practice, chunk size is a trade-off, and the optimal balance depends on the application domain and the LLM’s context window.

The query/retrieval workflow (online process)

The query workflow (Figure 4-3, part B) is an online process that happens at the real time when customers interact with Agent. When an agent receives a query, it first gets the query text’s embedding vector by using the embedding model. Then it searches for the vectors in the database that are closest to the query vector (using measures like cosine similarity). The agent retrieves the most relevant document chunks from the vector database by using the selected vectors, and ends by sending them to the LLM with the user query to generate an answer.

Why split text into small chunks?

LLMs have a limit on the maximum number of tokens they can process at once, including input and output tokens. This means there is limited space for context. That space is known as the context window. To accommodate these limitations, RAG segments the text it extracted into smaller, digestible chunks, so only a few highly relevant text chunks are sent to the LLM as context.

While Figure 4-3 shows a straightforward, “naive” RAG pipeline, production-grade RAG systems introduce additional complexity—such as chunking, ranking, deduplication, and multisource retrieval. For a deeper exploration of RAG, we recommend you start from Gao et al’s 2023 paper “Retrieval-Augmented Generation for Large Language Models: A Survey.”

Cache-Augmented Generation (CAG)

In the previous section, you saw how Retrieval-Augmented Generation (RAG) enhances language models by integrating external knowledge sources. While powerful, RAG also introduces several challenges.

There are extra steps and latency before calling the LLM. And when selecting relevant documents with limited tokens, RAG agents may select incorrectly. There is also additional complexity involved in building and maintaining a complicated search system.

As LLMs evolve, their context windows continue to expand. For instance, the Claude Sonnet 4 model (as of August 2025) supports a 1-million-token context window. With such capacity, the need to carefully select only a handful of “most relevant” documents diminishes. Instead, an entire knowledge base—or large portions of it—can be cached inside the model’s working context. This shift makes it possible to overcome LLMs’ key limitations (like static knowledge, hallucinations, and domain-specific gaps) without relying on a complex RAG system.

Cache-Augmented Generation (CAG) leverages the extended context windows of modern LLMs by preloading knowledge directly into the model’s KV cache. Rather than performing retrieval at query time, CAG loads relevant resources in advance so that the model can generate answers during inference using the cached context. This approach eliminates retrieval latency and reduces system complexity, while still enabling the LLM to ground its responses in external knowledge. Figure 4-4 compares RAG and CAG systems.

Figure 4-3. RAG and CAG comparison.

In RAG, every query triggers a search process to find context dynamically, while in CAG, once the knowledge docs are loaded into the LLM’s cache, queries are answered directly using the external knowledge in the context. This improvement greatly simplifies the agent design, as shown in Figure 4-5.

Figure 4-4. Agent workflow with CAG

Compare the agent design in Figure 4-5 to the one in Figure 4-2. CAG greatly simplifies the system architecture by removing the RAG component entirely. Instead of managing embeddings, indexes, and vector databases, the agent interacts directly with a CAG-enabled LLM already loaded with the knowledge and context. For a deeper dive into this paradigm, we recommend Chan et al’s 2024 paper “Don’t Do RAG: When Cache-Augmented Generation is All You Need for Knowledge Tasks.”

While CAG is promising, it also introduces its own challenges—particularly in terms of performance and efficiency. Large context windows and cache management can increase memory and compute requirements. In Chapter 7, we will explore techniques to optimize CAG model serving, with a focus on leveraging the LLM cache to improve efficiency.

How agents use model serving

As you saw from the knowledge agent example, autonomous agents often require orchestrating multiple models and tools to complete tasks. Typical components include:

LLMs for reasoning, planning, and dialogue generation

Embedding models for retrieval, similarity search, and semantic matching

Vision or speech models for multimodal perception

Tool-specific models for tasks such as code generation, classification, and summarization

External tools such as APIs, databases, or services, which the agent can invoke as needed

A key capability of modern agents is tool calling. Tool calling is a four-step process:

The LLM reasons over the user’s request and the available tools.

It generates structured output (often in JSON) that encodes the required inputs for the selected tool.

The agent executes the tool call.

The results are passed back into the LLM, which iterates until the task is complete.

This pattern allows agents to extend beyond pure text generation, combining reasoning with precise tool execution.

All of these models and tools are typically accessed through model-serving services (such as HTTP or gRPC APIs), so the agent can invoke them on demand. Since agents operate interactively and often in real time, high-performance, low-latency, and cost-efficient serving is critical to the success of any agent application.

Now that you’ve seen how models and tools are used in agentic applications, the remainder of this chapter will explore different approaches to building a production-grade model serving system.

LLM serving in enterprise systems: An overview

In this section, we introduce an abstract layered architecture that has been adopted by leading providers, such as OpenAI, Anthropic, and Together AI, as well as large enterprises, when deploying LLMs at cloud scale.

Compared with the simplified serving examples in earlier chapters, real-world deployments involve far more than simply hosting and executing models. At scale, providers must also address requirements such as authentication, pricing, resource management, networking, optimization, experimentation, observability, and on-call support. Meeting these business and operational demands requires multidisciplinary engineering teams of hundreds or even thousands of engineers.

Architecting such a system is challenging. The serving platform must enable many teams, each with distinct responsibilities, to contribute to a single evolving system without creating bottlenecks or excessive interdependencies. It must balance fast iteration with stability, all while operating under the constraints of cost, reliability, and user experience.

Figure 4-6 presents what a layered enterprise model-serving architecture looks like in practice. This conceptual view highlights the distinct requirements and challenges that arise at each layer.

Figure 4-5. An enterprise-level model-serving system architecture

Let’s walk through the layers of Figure 4-6 from top to bottom. At each layer, we’ll examine its purpose and key challenges.

Public API Layer

The public API is the primary external interface through which customers, developers, and internal services interact with the LLM system. It manages networking, authentication, pricing, rate limiting, and request routing.

Key challenges include:

High concurrency

This layer needs to support millions of simultaneous client connections.

Fair usage and monetization

This layer must enforce quotas, prevent abuse, and support accurate billing.

Low-latency global access

It uses geo-routing and regional caching to serve requests from the nearest region.

Security

This layer must ensure authentication, tenant isolation, and protection against network attacks.

Resource Management Layer

This layer manages the infrastructure hardware resources (such as CPU, GPU, memory, disk, and networking) required to serve models, use cases, and customer workloads across regions. Budget reviews and cost allocation are often integrated here.

Key challenges include:

Capacity planning

This layer must forecast demand while avoiding costly overprovisioning.

GPU utilization

This layer must manage heterogeneous GPU pools across datacenters and maintain high utilization.

Customer prioritization

This layer must enforce quotas and reservations for critical workloads, while enabling preemptible scheduling for lower-priority tasks.

Model Selection and Orchestration Layer

The orchestration layer determines which model(s) to use for each request, balancing accuracy, latency, and cost. Beyond selecting a single model, this layer can orchestrate multiple models in tandem, for example by applying speculative decoding or routing between model families.

Speculative Decoding

Speculative decoding accelerates inference by having a smaller “draft” model predict several tokens ahead, while a larger “target” model verifies them in parallel. If the draft is correct, the tokens are accepted; otherwise, corrections are applied. This technique reduces token-by-token latency without degrading quality. We will explore speculative decoding in more detail in Chapter 7, and you can also refer to the paper Fast Inference from Transformers via Speculative Decoding” (Leviathan, Kalman, and Matias 2022).

Key challenges include:

Cost–quality trade-offs

Not all tasks require the largest models. For example, OpenAI often defaults to gpt-4o-mini unless the user explicitly requests a stronger model. Also, triaging which model to use means this layer must understand the query.

Load balancing

This layer must prevent bottlenecks by distributing traffic across model pools.

Latency-sensitive use cases

For these cases, this layer can employ smaller or faster models or advanced optimizations such as speculative decoding.

Distributed Serving Layer

This layer sets up the distributed infrastructure for LLM model execution. This generally falls into two groups: distributed model hosting for large models, and distributed caching (such as KV caching, prompt caching, and semantic caching) to reduce unnecessary computation.

Here are some challenges addressed in this layer:

Hardware limitations

LLM sizes often exceed single-GPU memory.

Multi-GPU and multi-node coordination

Multi-GPU and multi-node distributed models must efficiently share the request’s context to maintain the serving service level agreement(SLA).

Caching for efficiency

This layer must route traffic based on cache content, such as KV cache aware routing, and reduce redundant inference for repeated or similar inputs.

Core Inference Layer

This is where models are actually executed. It integrates model-serving frameworks (vLLM, Triton, TensorRT-LLM, SGLang) with optimized kernels (FlashAttention, GEMM, PagedAttention) and exposes them as web endpoints. We explored this layer in detail in Chapter 3.

Model Optimization Layer

This layer focuses on applying all kinds of model optimization techniques to improve the model’s performance and efficiency without retraining from scratch. We will discuss different optimization techniques and their trade-offs in detail in Chapters 5, 6, 7, and 9.

Model Layer

The bottom layer provides the actual trained models to the serving system. It handles:

Moving models from internal training pipelines or external sources into production.

Organizing models by capability (voice, reasoning, video) and purpose (sandbox, experimentation, production).

Tracking and versioning models as the model evolves.

As the field advances, new models and technologies will reshape system design. However, the layered architecture pattern—separating concerns such as API, orchestration, inference, and optimization—remains essential. It allows teams with different responsibilities to innovate and optimize independently while contributing to a cohesive serving platform.

Building with an open-source stack

In this section, we discuss how to implement an enterprise model-serving system using a set of our favorite open-source software components. The goal is not to prescribe a single solution, but to illustrate practical design choices and give you a starting point for building your own serving platform.

First, figure 4-7 provides a quick system overview.

Figure 4-6. Implementing model serving with an open-source stack.

In this design, we chose to build an entire serving stack on top of Kubernetes, an open-source platform for automating the deployment, scaling, and management of containerized applications. Beyond its core features, Kubernetes also has a large open source ecosystem built on top of it, encompassing functions such as metrics, logging, networking, authorization and hardware management. Because of these, Kubernetes has been widely adopted and is the backbone for many successful cloud providers (Amazon AWS, Google GCP and Microsoft Azure) and foundation model vendors (OpenAI, Anthropic).

In figure 4-7, we utilize Kubernetes and its ecosystems to take care of resource management, networking, traffic routing and load balancing, service hosting and scaling, and service metrics and monitoring. To learn more about Kubernetes, we recommend the books Kubernetes: Up and Running by Brendan Burns et al. (third edition, O’Reilly, 2022) and Kubernetes Best Practices by Brendan Burns et al. (second edition, O’Reilly, 2023).

Implementing Public API

For our public API example, we decided to use FastAPI and Kubernetes to implement the web API, quota limits, auth, and tenant management. Let’s look at a few code examples.

First, using FastAPI, we built a sample web service (enterprise-model-api) that exposes a chat API (/v1/chat/completions). We also leveraged FastAPI’s dependency inject function (Depends) to attach authorization logic to this API:

@app.post("/v1/chat/completions")
async def chat(req: ChatReq, idp=Depends(require_auth)):
    await rate_limit(idp["tenant"])
    # select model, route traffic based on tenant information

For authorization, we had the option of using either a JWT token or an API key to authenticate and authorize user requests. During this process, the platform also retrieves customer metadata, which is essential for downstream steps such as quota enforcement, traffic routing, and model selection:

# authorize request either with jwt or api_key
async def require_auth(
   api_key=Depends(verify_api_key),
   claims=Depends(verify_jwt)):
    if not api_key and not claims:
        raise HTTPException(401, "Missing API key or JWT")
    tenant = claims.get("tenant") if claims else await \
        rds.hget(f"keys:{api_key}", "tenant")
    if not tenant: raise HTTPException(403, "Unknown tenant")
    return {"tenant": tenant, "claims": claims, "api_key": api_key}

async def verify_jwt(authorization: str):
    typ, token = authorization.split(" ", 1)
    if typ.lower() != "bearer":
        raise HTTPException(401, "Bad auth type")

    # verify JWT and return token content
    .. .. ..
    return jwt.decode(token, \
       jwt.algorithms.RSAAlgorithm.from_jwk(json.dumps(key)), \
       algorithms=[headers["alg"]], audience="api://llm", \
       options={"verify_exp": True})

def verify_api_key(x_api_key: str | None = Header(default=None)):
    ... ...

For autoscaling, we leveraged Kubernetes. In the following configuration, if the enterprise-model-api service’s average CPU usage exceeds 70%, Kubernetes will trigger horizontal scaling to add 3 to 15 more instances to this service:

apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata: { name: enterprise-model-api-hpa }
spec:
  scaleTargetRef: { apiVersion: apps/v1,
      kind: Deployment, name: enterprise-model-api }
  minReplicas: 3
  maxReplicas: 15
  metrics:
  - type: Resource
    resource: { name: cpu, target: { type: Utilization,
      averageUtilization: 70 } }

For rate-limiting requests to the public API service, Nginx and Envoy are the two most common approaches to use as ingress proxies. The following example shows a simple Kubernetes configuration using Nginx to limit traffic to the enterprise-model-api service:

apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
  name: api-ingress
  annotations:
    kubernetes.io/ingress.class: nginx
    # rate limit the traffic to 50 request per second max
    nginx.ingress.kubernetes.io/limit-rps: "50"
    nginx.ingress.kubernetes.io/limit-burst-multiplier: "5"
    nginx.ingress.kubernetes.io/proxy-body-size: "8m"
spec:
  rules:
  - host: api.yourorg.example
    http:
      paths:
      - path: /
        pathType: Prefix
        backend: { service: { name: enterprise-model-api, port: { number: 80 } } }

If you prefer Envoy, refer to the tutorial Enabling Rate Limits Using Envoy for step-by-step guidance. Both Nginx and Envoy integrate seamlessly with Kubernetes and can be enabled through configuration, highlighting Kubernetes’ strength and flexibility as a serving platform.

Implementing Model Selection

After the public API, let’s look at the model selection phase. Depending on the requirement, we can implement the model selection logic in a public API, in a separate service (middleware), or simply in a static routing config.

In our simplified example, the model selection logic is implemented within the public API service’s chat endpoint (/v1/chat/completions). Based on the token size of the incoming request, the system either:

Directly forwards the request to the selected model backend, or

Uses speculative decoding, where a faster, cheaper model generates multiple future tokens, and a slower but more accurate model verifies them in bulk. This approach avoids redundant step-by-step generation and provides significant latency improvements. (We will explore speculative decoding in greater detail in Chapter 7.)

Here is the logic:

@app.post("/v1/chat/completions")
async def chat(req: ChatReq, idp=Depends(require_auth)):
    await rate_limit(idp["tenant"])
    # choose the right model for the given request
    ep = choose_endpoint(req.model, idp["tenant"])

    # basic classifier: use speculation for long outputs, else direct
    if req.model.draft_enabled and req.max_new_tokens > 1024:
        draft_ep = config.get_draft_endpoint(req.model)
        # use draft model for generation, target model for validation
        gen = speculative_decode(req,
                 endpoint_draft=draft_ep,
                 endpoint_target=ep)
    else: # simple pass-through stream
         gen = passthrough(ep)

Speculative decoding has many forms of implementation

Speculative decoding has multiple variations, such as per-model adaptive speculative decoding and multi-draft decoding. While they all share the same foundational idea—using a smaller, faster model to reduce latency—they can be implemented at different layers of the architecture shown in Figure 4-6, depending on your specific optimization goals.

Now, let’s look at an example of self-implemented model endpoint-selection logic. The system first applies general preferences—such as region, cost, and load-balancing considerations—then layers on tenant-specific (customer) requirements to customize routing for individual clients:

def choose_endpoint(model: str, tenant: str):
    cfg = load_routes()
    # policy example: tenant allow-list, cost class, region, canary
    route = cfg["models"].get(model) or cfg["aliases"].get(model)
    if not route: raise HTTPException(404, f"Unknown model {model}")
    # weighted canary
    if "canary" in route and random() < float(route["canary"]["weight"]):
        return route["canary"]["url"]
    # tenant override
    route_over = (route.get("tenants") or {}).get(tenant)
    if route_over: return route_over["url"]
    return route["url"]

Implementing a model-serving endpoint

At the model hosting layer, this is where single-model and multi-model serving come into play. To implement this yourself, you can follow the designs introduced earlier in Figure 3-7 (Single Model Serving General Design) and Figure 3-8 (Multi-Model Serving Service Architecture), and build with tools such as vLLM, NVIDIA Triton, Kubernetes, and FastAPI.

If managing scalability and networking for serving instances feels too heavy, you can instead use Ray Serve—a scalable, framework-agnostic model serving library—for building online inference APIs. In the next few sections, we will demonstrate how to use Ray Serve to implement single-model and multimodel hosting.

Single Model Hosting on Ray

First, create a Python class called QwenVLLM to contain the serving logic. Let’s start by initializing the model (Qwen3) in vLLM:

class QwenVLLM:
    def __init__(self):
        # Initialize vLLM async engine once per replica
        args = AsyncEngineArgs(
            model="Qwen3-Thinking-2507",
            tensor_parallel_size=1,
            trust_remote_code=True,
            max_num_batched_tokens=4096,
            gpu_memory_utilization=0.9,
        )
        self.engine = AsyncLLMEngine.from_engine_args(args)

Then define the model serving request’s entry point and execution logic:

class QwenVLLM:
    # entry point for the generation/prediction request.
    async def __call__(self, http_request: Request) -> str:
        """
        Accepts a raw JSON body:
          { "text": "Explain KV cache simply.",
         "max_new_tokens": 200, "temperature": 0.7 }
        Returns the generated string.
        """
        body = await http_request.json()
        text = body["text"]
        max_new = body.get("max_new_tokens", 256)
        temp = body.get("temperature", 0.7)
        return await self.generate_text(text, \
            max_new_tokens=max_new, temperature=temp)

      async def generate_text(self, text: str, \
          max_new_tokens: int = 256, \
          temperature: float = 0.7) -> str:
          ... ...
        # run vLLM engine to execute Qwen model
        async for out in self.engine.generate(prompt, \
            params, request_id=req_id):
            # For vLLM with stream=True, incremental delta is here:
            delta = out.outputs[0].text_delta
            if delta:
                chunks.append(delta)
        return "".join(chunks).strip()

After implementing the model-serving code, you can host it in Ray Serve by adding a decorator @serve.deployment to the QwenVLLM class. With this declaration in place, during service deployment, Ray will automatically provision compute resources and deploy the service for Qwen model serving.

The following code instructs Ray to serve the Qwen3 model with three service instances, each allocated a dedicated GPU:

# Create RayServe instance for hosting Qwen3 model
@serve.deployment(
    name="qwen_vllm",
    # declare the number of service instances for the model
    num_replicas=3,
    # needs a GPU for speed
    ray_actor_options={"num_cpus": 0.5, "num_gpus": 1},
)
class QwenVLLM:
 .. .. ..

Multimodel hosting on Ray

If you need to host a group of models with similar input types on a shared service, Ray Serve offers a feature called model multiplexing to serve multiple models from a pool of service instances (replicas).

By using serve.multiplexed and serve.get_multiplexed_model_id API, you can extend your single-model hosting service to multimodel hosting in Ray. We’ll show you a quick example for hosting three different sentiment models on Ray Serve with two service instances.

First, you’ll define the model execution logic, which loads and executes the model in real time per your inference request:

@serve.deployment(
    num_replicas=2,  # two instances/replicas
    ray_actor_options={"num_cpus": 1} # resources per replica
)
class MultiModelService:
   # entry point for model inference requests
   async def __call__(self, request: Request):
       # Pick the model id for this request
       model_id = _serve.get_multiplexed_model_id()

       # load model and run model inference
       tokenizer, model = await self.load_model(model_id)
       inputs = tokenizer(text, return_tensors="pt")
       with torch.no_grad():
           logits = model(**inputs).logits
           probs = torch.softmax(logits, dim=-1).tolist()[0]

       return JSONResponse({
           "model_id": model_id,
           "label_probs": probs,  # [neg, pos] for SST-2 style models
       })

From the __call__ function of the MultiModelService class, you can see that the multimodel inference code’s implementation is very similar to that for single model serving. The main difference is that, in single-model serving, you preload the model in the __init__ function, while for multimodel hosting, you’ll try to load the requested model in real time for every prediction request using await self.load_model(model_id), since requests may ask for different models.

The drawback of this approach is that loading a model for every request introduces additional latency. To mitigate this, annotate the model loading function (load_model) with @serve.multiplexed. This tells Ray Serve to apply its caching mechanism when possible and to manage model swapping when resources such as memory are constrained.

The following examples illustrate how to enhance the model-loading logic with multiplexing enabled:

# Example registry: map model_id -> HF repo
# This model source could be S3 or any other storage
MODEL_REGISTRY = {
   "distilbert-sst2": "distilbert-base-uncased-finetuned-sst-2-english",
   "bert-sst2": "textattack/bert-base-uncased-SST-2",
   "roberta-sst2": "roberta-base-openai-detector",
}

@serve.deployment(ray_actor_options={"num_cpus": 1})
class MultiModelService:

   # enable multi model support and
   # set model replica count for each model
   @serve.multiplexed(max_num_models_per_replica=2)
   async def load_model(self, model_id: str) \
           -> Tuple[AutoTokenizer, torch.nn.Module]:
       """
       Lazily load & cache a model per model_id on this replica.
       Ray Serve will LRU-evict when cached models
         exceed max_num_models_per_replica.
       """
       repo = MODEL_REGISTRY[model_id]

       tokenizer = await asyncio.to_thread(\
              AutoTokenizer.from_pretrained, repo)
       model = await asyncio.to_thread(\
              AutoModelForSequenceClassification.from_pretrained,\
              repo)
       model.eval()
       return tokenizer, model

You can use a HTTP header ray_serve_multiplexed_model_id to specify the ID of the model to call in the prediction request. Ray will handle traffic routing, load balancing, and resource management on the backend. For example, the following curl request sends a sentiment prediction request to the distilbert-sst2 model host in RayServe:

curl -s -X POST "http://127.0.0.1:8000/SentimentService" \
  -H 'content-type: application/json' \
  -H 'ray_serve_multiplexed_model_id: distilbert-sst2' \
  -d '{"text":"I absolutely loved this!"}'

Building your own serving stack is not difficult

While the system may appear to involve many moving pieces, most of the components we’ve introduced in this chapter have already been battle-tested in production environments. By leveraging Kubernetes for infrastructure management, Ray Serve or Triton Inference Server for model hosting, and vLLM for distributed LLM serving and optimizations (like kernel-level acceleration and caching), you can assemble a robust, flexible serving stack.

With some lightweight custom logic in a self-implemented middleware microservice, it is entirely possible to build a full enterprise-grade serving solution yourself. Since the libraries and frameworks introduced here have matured over years of adoption, there are many strong reference implementations from which you can draw inspiration. In addition, AI coding assistants can significantly accelerate development by handling boilerplate and integration tasks.

Building with a Cloud Vendor

In this section, we use AWS SageMaker, Amazon’s AI platform, to illustrate six different ways to build a model-serving system on a public cloud. We will walk you through these options, from the fully managed end of the spectrum to the most customizable approaches. The goal is not to teach you how to operate SageMaker step by step, but to help you understand the fundamental logic behind cloud vendors’ serving options so you’ll be better equipped to navigate the landscape and make informed choices, even when the variety of options feels overwhelming.

Note

While the examples and options we list in this chapter are specific to AWS, their functionality is quite general. You can find corresponding components in other popular cloud vendors, such as Google GCP and Microsoft Azure.

In this section, we use the terms Docker image and Docker container interchangeably, both referring specifically to the model-serving service and its execution environment.

Option 1: Fully Managed Foundation-Model Serving

Amazon Bedrock is a fully managed service for serving foundation models via a simple API. Of the options we present here, it is the least customizable but the easiest to use. With Bedrock, you don’t handle training or hosting the model yourself. You simply select a model Bedrock supports (such as Amazon Titan, Anthropic Claude, or Stability AI’s Stable Diffusion), then make API calls to get predictions.

Bedrock provides several foundation models–both from AWS and from third-party model providers such as Anthropic and Cohere–that you can use out-of-the-box, without having to deploy or manage any infrastructure. This is ideal for quickly integrating AI capabilities into applications or prototyping with minimal setup.

The price model of Bedrock is “pay-as-you-go,” typically charging per request or per output token rather than, for instance, by the hour. This means you don’t run any servers in your account for these models – it means that AWS manages the scaling and computing resources internally.

As a result, Bedrock offers zero DevOps overhead: there are no servers to provision or scale and no model containers to build. The trade-off is that you are limited to the models and configurations AWS provides, with relatively limited customizability. For example, you cannot change the model’s architecture or training; you can only do moderate fine-tuning or prompt customization.

Using Bedrock just takes three steps:

Create an API key in your AWS account

Select the model you want to use from Bedrock’s model catalog

Initialize a Bedrock client and submit API calls

The following example illustrates this process:

# Set the API key as an environment variable
os.environ['AWS_BEARER_TOKEN_BEDROCK'] = "..."

# Create the Bedrock client
client = boto3.client(
   service_name="bedrock-runtime",
   region_name="us-west-2"
)

# Define the model and message
model_id = "us.anthropic.claude-3-5-sonnet-20240620-v1:0"
messages = [{"role": "user", "content": [{"text": "Hello! Can you tell me about Amazon Bedrock?"}]}]

# Make the API call
response = client.converse(
   modelId=model_id,
   messages=messages,
)

Limitation

Under the hood, what Amazon Bedrock provides is a serverless managed inference environment for foundation models. You interact with Bedrock via a high-level API, either through AWS SDKs or HTTPS endpoints.

When you call the Bedrock API to invoke a model, AWS’s infrastructure (which you do not see or manage) handles provisioning the GPU instances, loading the model weights, running the inference, and returning the result. This is conceptually similar to calling an external AI API (like OpenAI’s) – but it’s an AWS service call within your AWS environment.

Because Bedrock is fully managed by AWS, you cannot choose the instance type or customize the container or code. Nor can you deploy new models to Bedrock. If you have a proprietary model or a fine-tuned version of a model that isn’t provided, you cannot serve it via Bedrock; read the following sections to learn to host your custom model.

When to use

If you need immediate inference from high-quality foundation models and want the simplest integration possible, a fully managed solution like Bedrock is a great choice. It shines for use cases like building a quick prototype of a chatbot, text generator, or image generator using pretrained models, especially when you don’t want to manage any infrastructure.

However, it’s not suitable if you need to deploy your own custom model or if you require extensive custom inference logic. In such cases, other serving options provide more flexibility.

Option 2: One-Click Foundation-Model Deployment

Another no-code/low-code model-hosting method is Amazon SageMaker JumpStart. Like AWS Bedrock, SageMaker Jumpstart provides a model hub with pretrained models and easy deployment mechanisms. It also gives you a large catalog of models, including AWS-curated foundation models like Cohere, Falcon, Llama, and Stable Diffusion, as well as models from the Hugging Face Hub.

The big difference between JumpStart and Bedrock is that JumpStart lets you deploy these preselected models to your own SageMaker infrastructure with minimal coding. It’s often described as one-click deployment for models, and it’s available through the SageMaker Studio UI as well as programmatically via the SageMaker SDK.

If you’re hosting models in your own AWS account, you’ll pay by the hour for the SageMaker-managed instances. This means a bit more infrastructure management, since you now host a model endpoint yourself. You can choose the server instance types, such as g5.48xlarge and g6e.48xlarge,

In the example that follows, you can configure a few more things about the infrastructure than you could in Bedrock, such as the number of model-serving instances, the logging levels, and the server instance type:

# Define the model ID and version
model_id = "huggingface-llm-mistral-7b-instruct"
version = "*"  # Use the latest version

# Create a JumpStartModel instance with hardware configurations
model = JumpStartModel(
   model_id=model_id,
   model_version=version,
   instance_type="ml.g5.2xlarge",  # Specify the instance type
   role=sagemaker_execution_role, # Specify the execution role
   env={ # Example environment variable
       "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
       "SAGEMAKER_CONTAINER_LOG_LEVEL": "20"
   },
   # Add other configurations as needed.
)

# Deploy the jumpstart model
predictor = model.deploy(
   initial_instance_count=1,
   endpoint_name="my-mistral-endpoint",
   role=sagemaker_execution_role,
   sagemaker_session=sess
)

# Send prediction request
response = predictor.predict({"inputs": "Hello, world!"})

Limitation

Under the hood, SageMaker JumpStart uses the standard SageMaker inference infrastructure, but automates much of the setup. It provisions the server instances, downloads model artifacts and Docker images, and hosts the model as a model endpoint. Once deployed, the endpoint works like any other SageMaker endpoint: you pay for the uptime of the instance.

JumpStart is easy to use, but comes with some limitations on customization, Since it automates nearly every step of model deployment, you have limited knobs you can tune to customize your model. For example:

Not every model in the catalog supports fine-tuning or evaluation; some are deploy-only.

Instance choices are constrained per model; models have their own supported instance types.

You cannot customize the default model inference behavior, such as data pre-process and batching.

You can’t adjust prediction requests’ payload schemas or content types.

You can’t choose your CUDA or Pytorch version.

It’s difficult to apply optimization techniques since JumpStart configures serving options, such as context length and batch size, for the selected model.

If you need complete control over the inference logic, or if you need to serve a model that JumpStart doesn’t support, continue on to learn about more customizable SageMaker options.

When to use

If you want to get a popular pretrained model up and running in your own AWS environment quickly, JumpStart is ideal.

JumpStart targets no-code or low-code users who want to avoid dealing with container setup or writing inference code for well-known models. For instance, if you need a BERT model for text classification or Stable Diffusion for image generation and don’t want to manage the serving stack, JumpStart can deploy it for you. It’s also useful as a learning tool for SageMaker, since it sets up endpoints following best practices.

Option 3: Bring Your Own Model

For a still more customizable option to serve your own model or models that are not listed in JumpStart’s or Bedrock’s model catalogs, AWS SageMaker provides a variety of prebuilt serving Docker images called Deep Learning Containers (DLCs) that you can use to deploy models without writing custom inference code.

DLCs include framework-specific containers such as TensorFlow, PyTorch and Hugging Face Transformers, as well as SageMaker’s built-in algorithm containers, such as the DJL container. Think of this as a “bring-your-own-model” approach, but not “bring-your-own-code.”

With DLC, you can serve a trained model artifact (perhaps one you trained yourself or obtained elsewhere) without worrying about model-serving code, since the DLC knows how to handle that model type. You can use this for custom models (such as your own PyTorch model) or open source foundation models (for example, deploying a model from Hugging Face Hub) with minimal coding work.

The key difference from JumpStart is that, here, you are explicitly choosing which container (serving framework) to use and managing the deployment steps yourself, though still in just a few lines of code. JumpStart would autoselect those for you. This gives you a bit more flexibility – for example, you can specify exactly which version of PyTorch or Transformers library you want, or deploy a model that isn’t in the curated JumpStart list. However, to use this option, you need to understand the container’s requirements.

Let’s walk through an example of hosting a PyTorch model with TorchServe. The first step is to select a prebuilt DLC that matches your framework and Python version. You can use the following code to search for a suitable DLC image for your PyTorch model.

# Find the available serving image by searching model framework
# and server instance type
baseimage = sagemaker.image_uris.retrieve(
        framework="pytorch",
        region="",
        py_version="py310",
        image_scope="inference",
        version="2.0.1",
        instance_type="ml.g4dn.16xlarge",
    )

Next, create a serving image and a Model object with model file location in AWS S3. Deploy the model object to the selected instance type (for instance, g4dn.16xlarge). The model.deploy function will handle resource provision and service deployment on AWS.

# Create the Model object
model = Model(model_data = f'{output_path}/mnist.tar.gz',
                  image_uri = baseimage,
                  predictor_cls = Predictor,
                  name = "mnist")

# Deploy model
endpoint_name = 'torchserve-endpoint-1'
predictor = model.deploy(
    instance_type='ml.g4dn.16xlarge',
    initial_instance_count=1,
    endpoint_name = endpoint_name,
    serializer=JSONSerializer(),
    deserializer=JSONDeserializer())

Here’s another example: serving an open-source LLM (Llama) from Hugging Face using AWS’s prebuilt Large Model Inference (LMI) containers, powered by the DJL serving framework with vLLM.

# Create the SageMaker Model object.
# In this example we let LMI to configure the deployment settings
# based on the model architecture
model = DJLModel(
  model_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
  env={
    "HF_TOKEN": "",
    # Add more serving configurations here
    # Example: set tensor parallel degree
    "OPTION_TENSOR_PARALLEL_DEGREE": "4",
    # Example: specify serving loader as vLLM
    "OPTION_SERVING_LOADER": "vllm",
    # Example: set max rolling batch size
    "OPTION_MAX_ROLLING_BATCH_SIZE": "128",
  }
)

# Deploy your model to a SageMaker Endpoint
# and create a Predictor to make inference requests
endpoint_name = sagemaker.utils.name_from_base("llama-8b-endpoint")
predictor = model.deploy(
    instance_type="ml.g5.12xlarge",
    initial_instance_count=1,
    endpoint_name=endpoint_name)

This approach requires a more explicit setup than JumpStart does: you need to search for and specify the appropriate DLC container and version to match your model’s framework, and ensure your model artifact is properly prepared. In return, you gain flexibility—for example, you can newer releases of frameworks not available through JumpStart’s simplified interface. You still don’t need to write serving code yourself, since the container provides the implementation.

Limitation

The main limitation of the DLC approach is that you have little control over the container internals. Frameworks, NVIDIA drivers, CUDA, Python, and the base OS are all pinned to the image tag (for example, TensorFlow 2.19 with CUDA 12.2 on Ubuntu 22.04). To mix versions of the internal libraries, you must build your own image.

SageMaker also enforces a fixed port and HTTP contract: each DLC container must expose a web server on port 8080 and implement /invocations (POST) and /ping (GET). These cannot be changed without leaving the standard path, and default response timeouts apply (for example, 60 seconds for standard responses).

Another constraint is that DLCs assume specific input/output formats for prediction requests. If your use case requires custom preprocessing or postprocessing, you’ll need to either provide a custom inference script (Option 4) or handle those steps outside the model.

When to use

This prebuilt serving container (DLC) approach is ideal when you have a model trained with a common framework, you want to deploy it quickly using SageMaker’s managed service, and you might need slightly customized choices. For example:

If you’ve fine-tuned a Hugging Face Transformers model (like BERT or Qwen3) and have the weights, you can use the Hugging Face inference DLC to serve it with minimal effort.

If you want to deploy a TensorFlow SavedModel or a PyTorch model.pth file, you can use images like TorchServe or Tensorflow Model Serving to serve it with minimal effort.

To use the DLC prebuilt serving images approach, your model must be in a supported format and you must be okay with the container’s default prediction handling. If you need to preprocess input or postprocess output in non-AWS-standard ways, you might eventually require the next level (custom inference script, option 4). But if the default serving logic is sufficient, this saves you from writing any inference code.

Option 4: Bring Your Own Code

As we move up the customization ladder, the next step is to write your own serving code. In AWS SageMaker, this is done through Script Mode for inference. With Script Mode, you still rely on a SageMaker-provided container for the underlying framework, but you supply an entry-point script that implements the logic for model loading and prediction. This gives you much greater flexibility: you can define custom preprocessing and postprocessing, support nonstandard input/output formats, and even load multiple models within a single container. The trade-off is that you now need to write and maintain code, whereas the earlier options required no custom inference code.

As an example, let’s use a SageMaker Large Model Inference (LMI) container to host an LLM with our own serving implementation. For brevity, we’ll illustrate this concept with pseudocode.

The first step is to define the per-model serving configuration. The following code specifies both the serving engine and location of the model.

%%writefile my-own-llm/serving.properties
engine=Python
option.tensor_parallel_degree=2
option.rolling_batch=vllm
option.s3url=s3://sagemaker-us-west-2-/large-model-lmi/code/my-own-llm

Next, we implement our own model serving (loading and execution) code in a file called model.py and implement the model-serving execution entry point at the handle function:

%%writefile my-own-llm/model.py
import os
# ...
PAD_TOKEN_ID = 50256

# initialize model
def initialize(properties):

    model = AutoModel.from_pretrained(model_id)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    .. ..

# execute model
def run_inference(input_texts, onnx_model, tokenizer):

    max_batch_size = 128
    z = torch.empty([0,768]).to("cuda")
    for i in range(0, len(input_texts), max_batch_size):
        logging.info(f"Start Iteration: {i}")

        start_time = time.time()
        batch_dict = tokenizer(input_texts[i:i+max_batch_size],\
             max_length=512, padding=True, truncation=True, \
             return_tensors='pt').to("cuda")
        logging.info(f"After Tokenize Latency: {time.time() - start_time}")

        start_time = time.time()
        with torch.no_grad():
            outputs = model(**batch_dict)

        # ... ..
        z = torch.cat((embeddings,z), 0)
    results = [{"embedding": embedding.tolist(), "index": idx} \
        for idx, embedding in enumerate(z)]
    return {"embeddings": results}

# model execution entry point
def handle(inputs: Input) -> None:
    logging.info("handle : Handle start...")
    global model, tokenizer
    if not model:
        logging.info("handle : initializing model")
        model,tokenizer = initialize(inputs.get_properties())

    if inputs.is_empty():
        logging.info("handle : inputs is empty")
        # Model server makes an empty call to warmup the model on startup
        return None

    # implement model request pre-process logic
    data = inputs.get_as_json()
    input_sentences = data["inputs"]
    # run model inference
    res = run_inference(input_sentences, model, tokenizer)
    logging.info("handle : Handle End")
    return Output().add_as_json(res)

Next, we package mode.py, the model weights, serving.properties, and other model configuration files into one file and upload it to S3 cloud storage:

# package model file
%%sh
tar czvf my-own-llm.tar.gz my-own-llm/

# upload model file to S3
s3_code_prefix = "large-model-lmi/code/my-own-llm" # increment the version
bucket = sess.default_bucket()  # bucket to house artifacts
code_artifact = sess.upload_data("my-own-llm.tar.gz", bucket, s3_code_prefix)

Once we have the model file with the customized serving code and configuration, the last step is deployment, which is very similar to option 3, including choosing a serving image, creating a model object, and deploying the model.

# choose a LMI image.
image_uri = image_uris.retrieve(
    framework="djl-deepspeed",
    region=sess.boto_session.region_name,
    version="0.25.0"
)

# create model object
model = Model(image_uri=image_uri, model_data=code_artifact, role=role)

# deploy model to an endpoint
predictor = model.deploy(initial_instance_count=1,
             instance_type="ml.g5.2xlarge",
             endpoint_name="my-own-llm-128")

Limitation

In many cases, the “Bring Your Own Code” approach strikes a balance: you can tailor the inference logic to your specific model and use case while reusing the cloud vendor’s serving stack (AWS LMI or DLC containers) and infrastructure (SageMaker model deployment).

The trade-off lies in what you inherit from the vendor’s container. You’re limited to the frameworks, Python/OS/CUDA versions, and serving libraries baked into that image, and you cannot change the HTTP interface or request schema. If you need full control—such as using a preferred library like vLLM, customizing the runtime stack, or defining your own APIs—you’ll need to move to Option 5: Bring Your Own Serving Image.

When to use

Essentially, when your prebuilt containers’ limitations (see option 3) become a bottleneck, script mode lets you override the per-model serving logic in the DLC without building a container from scratch.

More specifically, use this option when you have a model that can run in one of the supported SageMaker containers, but the inference process needs custom steps or is not covered by the default container behavior: for example, if you need to transform the input data significantly before feeding it to the model (such as custom encoding and image transformations) but process the model output similarly (like filtering or formatting results).

Option 5: Bring Your Own Serving Image

At the top end of our first five options for customization and control, the most customized option is to bring your own Docker container to SageMaker. In this scenario, you are responsible for providing a container image that can serve your model.

SageMaker will still manage the deployment of this container to a cluster of instances and provide the endpoint URL, but everything inside the container is up to you. This is the most flexible option: you can use any programming language, any framework (even ones not directly supported by SageMaker’s built-in containers), and any custom code for inference. You also can install any system dependencies, use custom network configurations, and so forth within the container.

Essentially, SageMaker treats your container as a black box that adheres to a simple contract: HTTP requests come in, responses come out on certain ports and paths. This is where the single-model and multimodel serving services we introduced in chapter 3 fit in: You can put the service serving code into the Docker container and still use SageMaker to deploy and host the service and models.

Limitation

The main drawback of bringing your own serving image is the added complexity: you must build, maintain, and update the container yourself, while ensuring it complies with SageMaker’s requirements, such as its health check and invocation endpoints. Testing and debugging are also more involved—you’ll need to write a Dockerfile, implement the app, push the image to AWS ECR (Elastic Container Registry), and work with more verbose API calls. This overhead is the trade-off for having full control.

Another limitation is in distributed serving, where performance depends on coordination across serving instances. If you want to implement optimizations such as distributed KV caching, prompt caching, or request routing, it becomes necessary to move to Option 6.

When to use

In short, whenever the constraints of SageMaker’s provided containers become too limiting, “Bring Your Own Serving Image” is the escape hatch. With this option, you can customize virtually anything inside the container—for example, you could adopt a different ML framework or implement performance-critical logic in C++. As long as your container complies with SageMaker’s required HTTP interface, it will run.

Some common scenarios include:

Your model depends on a stack or serving framework that SageMaker doesn’t support, isn’t yet available in the latest version, or has been heavily customized.

You need complete control of the inference environment, beyond model code—such as system-level optimizations or running on a specific Linux distribution.

You want to bundle multiple processes or services in the same container for advanced use cases, like running an auxiliary service alongside the model to collect specialized metrics based on price or user metadata.

Option 6: Build Your Own Serving Infrastructure

The first five options all leverage parts of AWS’s serving stack. But if you want to operate a full-fledged model-serving platform—like the one outlined in the “LLM Serving Enterprise System Overview” section—you can build your own services directly on top of cloud infrastructure. For example, you could use a cloud-managed Kubernetes service.

In this model, you build custom serving images, choose your runtime (such as, Triton, vLLM, TensorRT-LLM, KServe, or Ray Serve), and run them on GPU node groups. You also take ownership of traffic management, autoscaling, security, observability, and cost controls. At the same time, you can still lean on AWS primitives and managed components, such as EKS, ALB/NLB, ECR, S3, EFS/EBS, IAM/IRSA, Karpenter/Cluster Autoscaler, and CloudWatch, a and you can use Kubernetes add-ons like Prometheus/Grafana, OpenTelemetry, Fluent Bit, CNI/NetworkPolicy, Gatekeeper/Kyverno, and ArgoCD/Rollouts.

When to use:

You require full control of the serving runtime, including kernels, batching, tokenization, custom sidecars, or nonstandard APIs such as gRPC or Server-Sent Events (SSE).

You need aggressive cost and performance tuning, such as spot GPU utilization, GPU sharing (MIG/MPS), model bin-packing, or custom autoscaling based on tokens/second.

You must meet strict compliance or data isolation requirements, including private clusters, VPC-only egress, per-tenant isolation, or custom audit trails.

You want to run on cutting-edge hardware or implement specialized optimizations not yet supported in managed services—for example, speculative decoding, KV-cache sharding, or advanced routing and load-balancing strategies.

Comparing the Options

There is no one-size-fits-all serving solution. Each company and each team has a different capacity, business model, and project timeline. When we evaluate which option will work best in a specific case, we normally start from the trade-off between ease of use and control, then evaluate from several other perspectives, such as cost and performance.

Figure 4-8 presents a decision tree to guide you in choosing among the available options, beginning with ease of use.

Figure 4-7. Decision tree for making serving choices based on ease of use

By following this decision tree, you can make serving choices based on ease of use—that is, how little development work and customization are required. The more AWSservices and components you leverage from, the faster you can get a model up and running.

In practice, however, development speed is only one factor. Operational and maintenance costs are equally important. Many teams start with the simplest option to validate an idea, then shift toward more customized solutions as projects mature or requirements grow.

For example, you might initially run the Qwen3 model on Bedrock, which charges at a rate of $0.10 dollars per input million tokens. As usage (throughput) increases, you might decide to move to Option 2 or 3 to host Qwen3 in your own account, paying $1.172 per hour for server instances, a reduction compared to the per-token cost. Later, as business needs demand higher throughput and lower latency, you could transition to options 4, 5, or 6, depending on the level of customization required.

Build or Buy? Understanding Strategies

We’ve heard lots of people debate whether to build their own serving stack or use services provided by cloud vendors. Both authors of this book have been building and maintaining different model-serving systems for more than a decade, and we can tell you that this is not binary. The choice between built and rent is a spectrum, not a switch!

The better you understand how to build a serving stack from first principles, the better you can evaluate managed options from cloud vendors, spot hidden trade-offs, and decide how much to customize at any moment. As with the decision tree from figure 4-8, think in terms of degrees of control, rather than either/or.

In practice, most teams operate somewhere in the middle: they run on vendor platforms while adding targeted customizations such as custom handlers, autoscaling signals, routing logic, or cost controls. A full “build your own” (BYO) approach is usually reserved for cases where control, large-scale cost efficiency, or strict compliance requirements make it necessary. For example, you might choose BYO to avoid vendor lock-in and ensure your system can run across multiple cloud providers.

Why knowing how to build helps—even if you won’t build

If you’re using a cloud vendor, even though you won’t need to build everything yourself, understanding a homegrown serving stack (as described in chapter 2 and 3) is critical. This understanding lets you:

Decode vendor features

You’ll recognize which knobs map to batching, caching, quantization, and adapter loading, and where the ceilings are.

Quantify trade-offs

You’ll be able to perform a good trade-off analysis between different vendor features and approaches, such as “bring your own model” versus “bring your own code.”

Customization

You can keep the vendor defaults for 80% of paths and surgically replace the 20% that need deeper control.

Debugging and troubleshooting

You can get things working or fixed quickly in the vendor system. Vendor solutions often work like black boxes, with limited documentation for their configuration and internal logic. Logging can be another pain point, since not all the logs are captured and displayed.

Our Selection Strategy

Here is an approach we’ve found very useful in practice, when we discuss with clients whether to stay with their vendor-provided serving solution or move to a more customized solution:

Stay vendor managed if your SLOs (such as prediction latency and throughput) are met, costs are acceptable, and feature velocity and timelines matter most.

Hybridize when a few model endpoints need special batching, routing, or performance tuning or per-tenant isolation.

Go BYO when you need control over hardware, runtime, or sovereign networking, or when the cost-of-serving demands deep tuning.

Go back to the vendor-managed approach if the volume stabilizes at a low rate or if the system’s complexity isn’t paying off.

When you learn to build a serving system, you can choose how much to build and customize at any point. The winning strategy is a dynamic one: keep a stable API and shared telemetry, customize where it pays off, and expect your position on the spectrum to change as your product, traffic, and constraints evolve.

Measuring Performance in LLM Serving

To prepare for the LLM optimization techniques we’ll show you in the next two chapters, we’ll now introduce you to two critical metrics we have observed in practice while operating the model-serving service. These quantitative measurements allow us to accurately and reliably assess the impact of each optimization. They are:

Latency

The time required to process and generate responses.

Throughput

The number of requests or tokens processed per unit of time.

Latency Metrics

When measuring the speed of LLM prediction execution, the three most important metrics are end-to-end latency, time-to-first-token (TTFT), and time-per-output-token (TPOT). Let’s look at them one by one:

End-to-end latency

E2E latency is a relatively general term used to measure the total computation time taken for both ML and non-ML workloads. In the context of LLMs, it refers to the time from when the model receives a request to when it completes generating the full response. E2E latency can also include the overall serving latency when needed, which accounts for not only model execution time but also additional delays (processing time) in the model serving system, including request queuing (when taking big loads of requests), network delays, routing time, scaling overhead, and other various system-level factors.

Time-to-First Token (TTFT)

TTFT refers to the time elapsed between the model receiving a request and emitting the first token. From users’ perspective, in chatbot-like use cases, it translates to how responsive the LLM model is–meaning how quickly the customer will see the first output token returned from the model. It also corresponds to the Prefill phase we discussed in Chapter 2, where the model processes the input prompt, attends to its full context, understands it, and prepares to generate output. In practice, TTFT is mainly used when discussing model latency.

Inter-Token Latency (ITL) or Time-per-Output-Token (TPOT)

ITL (or TPOT) measures the time taken to generate all subsequent tokens after the first. It corresponds to the Decode phase (see Chapter 2), where the model generates tokens sequentially, typically at a relatively steady rate. This metric is crucial for understanding the efficiency of a LLM’s autoregressive token generation. Like TTFT, ITL/TPOT is also mostly relevant when discussing model latency.

Figure 4-8. Key latency metrics in LLM execution.

Figure 4-9 shows how TTFT and ITL map to the Prefill and Decode phases within overall end-to-end latency in an LLM execution. The time costs of tokenization and detokenization are also captured in the graph.

In practice, tokenization and detokenization times are typically minimal compared to model inference time, and are often ignored from TTFT and ITL calculations

Now that you’re clear on the latency-related concepts, let’s look at the formulas to calculate them. We’ll assume that ITL is consistent; call the total number of tokens generated from the LLM N; and ignore tokenization and detokenization time:

The formula for E2E latency is:

E2E latency = TTFT + ITL x (N - 1)

To calculate TTFT, the formula is:

TTFT = full tokenization of the input + prefill + decode first token + de-tokenize first token

= prefill + decode first token

Finally, to calculate ITL:

ITL = decode one new token + detokenize the new token - detokenize the last token

= decode one new token

In real-world use cases, all three latency metrics—E2E Latency, TTFT, and ITL—are important, but their priority varies depending on the application and use case.

For example, if an LLM call is part of an agentic workflow where the downstream process depends on the full model output, then minimizing E2E latency becomes the primary objective. On the other hand, in chatbot applications, where responses are streamed back to users in real time, TTFT can have the most significant impact on user experience because it directly corresponds to the model’s responsiveness. But if the output is very long, ITL becomes a critical factor as well, since a very high intertoken latency can make the response feel sluggish.

In the following chapters, we will explore techniques for improving LLM serving latency. However, not all optimization techniques will enhance all three metrics at the same time. Certain trade-offs need to be made, so selecting the right techniques will depend on the specific requirements of the use case.

Throughput Metrics

Throughput is one of the core metrics in LLM serving. It’s used to answer the question: “How many predictions can the system handle over time?” This tells you how efficiently your serving system uses resources to deliver results at scale. The commonly used throughput metrics are:

Requests per second (RPS) or requests per minute (RPM)

RPS or RPM are universal throughput metrics for both ML and non-ML workloads. These metrics quantify how many requests the model or the entire serving infrastructure can handle within a given timeframe. But to LLM serving, RPS and RPM have quite a lot of drawbacks and lack of granularity specific to LLM serving, which we will discuss later.

Tokens per second (TPS)

TPS is a throughput metric specifically used for LLM generation. It measures the number of output tokens the model or the LLM serving service generates per second. One thing to keep in mind is that TPS only counts generated tokens: not input tokens or the sum of input and output tokens.

Both RPS/RPM and TPS have advantages and limitations, and neither is a perfect standalone measure of efficiency in LLM serving. Let’s look at why.

First, RPS and RPM are sensitive to prediction request patterns, since long inputs and outputs can negatively affect LLM serving performance if measured per request. As a result, RPS and RPM heavily depend on factors like input and output length, as well as the number of concurrent users. This means direct RPS/RPM comparisons between two different workloads with different traffic patterns are basically useless: it’s not an “apples to apples” comparison.

Second, TPS is useful, but it has risks because it can be manipulated. TPS is a standardized metric in the LLM space, and in a lot of cases we also leverage it to calculate cost per token. However, it can be artificially inflated in a few ways:

Reducing input length significantly lowers TTFT, which in turn reduces the per-request workload. This makes the TPS appear higher than it should.

Larger or optimized batch requests can also increase TPS. Increasing batch sizes to maximize model process parallelism or making the batched request have uniform input and/or output lengths can improve TPS, since doing so improves GPU utilization and reduces GPU idle time. We will explain why in the next 2 chapters.

Best Practices for Performance Measurement

When introducing performance measurement, we noted that no single metric is sufficient on its own. Latency and throughput are the two core metrics, but their value depends on how they’re interpreted and applied in context. In our experience, the most effective performance optimization comes from combining these metrics with a set of practical best practices. Let’s walk through some that we’ve found especially useful.

Identify the trade-off between latency and throughput

A good starting point is to recognize the trade-off between latency and throughput. Improving one often comes at the expense of the other. For example, for offline, noninteractive workloads such as batch processing, throughput usually matters most because it lowers cost. In contrast, interactive applications like chatbots or real-time agents demand responsiveness, making latency the critical metric.

Set acceptable latency goals based on specific use cases

The next important thing is defining what “good enough” means for your application. Setting clear latency targets prevents unnecessary optimization. For example, in an interactive chatbot, if you’ve already achieved a 1-second response time, squeezing it down to 0.5 seconds may not significantly improve user experience. At that point, it’s better to shift focus toward optimizing throughput and cost.

Decompose E2E latency into TTFT and ITL

Breaking down end-to-end latency into more specific components—time-to-first-token (TTFT) and inter-token latency (ITL)—provides a sharper view of performance bottlenecks. These sub-metrics can guide targeted improvements. For instance, in systems with capped output length, TTFT may be more critical than ITL, since users care more about how quickly the first token appears than how quickly the rest are generated.

Simulate real-world traffic patterns as closely as possible

Another key practice is simulating traffic patterns that resemble actual usage. Input and output lengths vary widely across users, and both heavily influence performance. When creating test suites, we need a clear understanding of the use case and its real-world traffic patterns to ensure performance optimizations are meaningful and do not deviate from real usage.

For example, long prompts with short answers behave very differently from short prompts with long outputs. If the setup matches the real usage, you will be optimizing in the correct direction as a whole and making the right decisions and trade-offs.

When simulating traffic, please remember that real-world traffic to systems is never perfectly uniform. Traffic levels can be different throughout the day or at specific times of the year. Bursts of traffic can overload the system, leading to queuing delays or even request failures. So please analyze your traffic pattern carefully before simulating it.

Maintain consistency in your experiments

Consistency across experiments is just as critical. Performance tuning often involves adjusting multiple parameters, and their interactions can be complex. Changing too many at once makes it impossible to isolate effects. What we found useful is to turn one knob at a time to isolate the impact of each optimization. Keeping request patterns and configurations stable across tests ensures reliable, apples-to-apples comparisons.

Monitoring and tracking hardware utilization

Monitoring hardware utilization provides another valuable perspective. Tracking GPU and CPU usage, memory consumption, and related metrics helps identify whether bottlenecks are caused by model behavior or hardware limits. This understanding is key to applying optimizations effectively, as later chapters will show.

Avoid manipulating or inflating performance metrics

Ensuring a fair and representative benchmark is important to prevent you from drawing misleading conclusions. As discussed in the TPS and RPS/RPM section, performance metrics can be artificially inflated and doing so will cause backfire later on when the system goes to production later on or lead the optimization effort to the wrong direction.

Continuously monitor performance in production

Once the model gets deployed, continuous monitoring ensures stable performance and detects possible shifts in usage patterns. Beyond latency, throughput, and hardware utilization, user behavior can also shift over time, causing request spikes and failures.

Run test suites periodically

Even after the model goes to production, regular testing is needed to prevent regressions, drive continuous improvements and keep the system efficient, scalable, and aligned with evolving workloads.

For example, regression checks can ensure that new updates don’t degrade latency or throughput.; Scalability tests can simulate peak loads on days like Black Friday or scale down for nonpeak hours. A/B testing can create data-driven, apples-to-apples comparisons on new optimization techniques you want to introduce before you roll them out to production.

In short, meaningful performance measurement requires a mix of granularity, realism, fairness, and ongoing vigilance. By applying these best practices, you’ll avoid misleading results and instead gain insights that translate into real-world improvements in latency, throughput, and cost efficiency.

Summary

In this chapter, we built on Chapters 2 and 3 (inference internals and single/multimodel design) and stepped you up to production best practices for LLM serving—agentic use, enterprise architecture, build-vs-cloud choices, and the metrics you’ll use to optimize serving performance.

We began with model serving in agentic applications, walking through a minimal Knowledge Agent that uses OpenAI models, a planner, and actions (RAG for retrieval, summarization, analysis) to show how serving enables autonomy. We then compared RAG and CAG (preloading knowledge into context/KV-cache) and analyzed their trade-offs.

Next, we presented a layered enterprise architecture whose components can evolve independently, plus an open-source Kubernetes stack to implement it. We also covered six AWS serving options spanning increasing customization and provided a decision tree for choosing among them.

A key theme is that the decision between building and using a cloud vendor is a spectrum, not a binary. Most teams operate in the middle—leveraging vendor platforms with targeted customizations—and shift that balance over time based on ease-of-development, cost, and usage patterns.

Finally, we standardized key metrics (latency, throughput) and best practices for fair benchmarking, realistic traffic simulation, hardware-utilization tracking, and regression/scalability testing. The next chapters will show you how to turn these metrics into concrete latency, throughput, and cost improvements.

Chapter 5. Challenges When Serving LLMs

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 5th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

So far in this book, we have demonstrated the core concepts of model serving, several architectural patterns for serving ML models, and the analysis of the trade-offs involved in deploying models at scale. By now, we hope you have a strong understanding of model-serving paradigms, because we are about to take a significant step into a different realm. In this chapter, we will shift our focus to one of the fastest-growing fields in the AI world: optimizing large language models (LLMs) for serving.

Since the rise of ChatGPT in late 2022, LLMs have transformed how AI is applied in real-world scenarios, from chatbots and code generation to advanced reasoning and decision-making systems. However, their sheer size, computational demands, and unique serving requirements introduce challenges that can go far beyond classic model-serving techniques. From novel ideas to widely adopted frameworks, the field of optimizing LLMs for faster and more efficient serving performance has evolved at an unprecedented pace. It can be daunting: anyone not familiar with this area can easily get overwhelmed. For example:

When reading technical blogs, you may wonder: “What is this vLLM framework that has gained popularity in just a year or two and has already been adopted in so many places?”

When reading research papers, you may ask: “How does FlashAttention work, and how can I optimize it at the hardware level to speed up LLM inference?”

When following AI news, you may come across a new term called MLA and ask: “How much more efficiently does the DeepSeek V3 model run than other attention mechanisms in the serving phase?”

In the following chapters, we will introduce you to all these advancements. But before we go into each technique, we need to establish a strong foundational understanding and gradually move you from basic concepts to more complex techniques. This chapter bridges between all the serving concepts, principles, and paradigms you’ve seen in the prior chapters and all the advanced LLM optimization techniques we will cover in the following chapters.

You first need to understand the landscape of LLM serving—why it is important, what the hardware requirements are, and how to build your intuition about model optimization. Specifically, we will cover why serving an LLM efficiently can be crucial to the success of your application and business at the very beginning. We’ll look at modern hardware’s role in LLM serving next, exploring AI accelerators such as GPUs to understand their intricacies of memory, compute capability, and interconnect functionality. Then we’ll cover:

the major bottlenecks in LLM serving and how to mitigate them

constraints when loading LLMs for serving

bottlenecks when executing LLMs, specifically in Prefill and Decode

why model serving can become bottlenecked in different phases

To discuss the last point, we introduce a new concept called arithmetic intensity.

Understanding these concepts is essential because, without the fundamentals, optimizing LLM serving can become a tiring trial-and-error experiment. It’s easy to get stuck in local optima without knowing it, or just knowing how but not why.

By the end of this chapter, you will have the intuition you need to navigate the world of LLM optimization with confidence. The next chapters will build upon this foundation by going through the various techniques needed to improve performance, reduce costs, and deploy LLMs efficiently for the workload your application requires in production in greater detail.

Why Optimizing LLM Serving is Important

The previous chapters have covered how to make ML models functionally operational and well-designed when deploying them to production. Another crucial aspect is ensuring that the ML model runs quickly and efficiently in a production environment. This is especially important for serving LLMs, which require substantial compute power and memory from hardware.

To better understand their impact, we categorize the key factors into three aspects:

Customer experience

Cost efficiency

Scalability, peak load handling, and feasibility

Customer experience

Customer experience is key to any product’s success and can be highly correlated with model response latency. Imagine asking a question in ChatGPT and having to wait over 20 seconds to receive the first token. Such a delay is unsatisfactory, and customers can easily become frustrated while waiting. If we can optimize the serving performance to reduce the model’s latency from 20 seconds to just 1 second, while keeping the same hardware setup and overall throughput, that would be a game-changer for the product’s success.

However, extremely fast response times may not always yield significant benefits. For example, if we reduce the time to receive the first token from 0.1 seconds to 0.01 seconds, the difference is so small that it doesn’t register in our human perception. In this case, if we can trade off some latency for higher throughput, increasing it from 0.01 seconds to 0.1 seconds, we will achieve better cost efficiency.

Figure 5-1. Model response latency vs. customer satisfaction

Figure 5-1 illustrates the relationship between latency vs. customer satisfaction. As you can see, the two are inversely correlated—when latency is high, customer satisfaction drops significantly, approaching zero. As latency decreases (moving from right to left along the X axis), customer satisfaction improves. However, as latency continues to decrease, the incremental gains in customer satisfaction become a lot less pronounced.

Another aspect of customer experience is generation quality. Within the same family or similar families of models, the larger models generally produce higher-quality responses. For example, without optimization, we might be limited to an 8-billion-parameter model to meet latency requirements. However, after extensively optimizing the whole model-serving system, we could achieve the same latency and throughput on the same hardware setup while leveraging a model with 32 billion or even 70 billion parameters. This can significantly improve response quality.

Figure 5-2. Model quality score comparison between Llama-3-70B and Llama-3-8B. (Source: Open LLM Leaderboard Model Comparator)

Figure 5-2 compares the benchmark scores ofLlama-3 8B and 70B, demonstrating how larger models outperform smaller models in the same family.

Cost efficiency

Beyond delivering a great customer experience, a product must also be financially viable. No matter how powerful or valuable an AI system is, it cannot succeed if it is too costly to operate at scale. Among all the costs of AI development and operation, model inference is the most expensive component. This might seem surprising, as training often appears to be the dominant cost—especially with reports stating that training costs for LLMs like GPT-4 exceeded $50 million.

However, as shown in Figure 5-3, in which Signal Integrity Journal projects AI-chip revenue trends from 2024 to 2034, inference hardware consumption already surpasses training today—and the gap is expected to widen even further. This reflects the growing demand for efficient inference infrastructure, as AI applications scale to serve millions of users in real time.

Figure 5-3. AI-chip revenue projected growth and composition change between training and inference. (Source: Signal Integrity Journal)

To better understand the dominance of model serving costs in Figure 5-3, let’s do some further analysis. Training costs are primarily upfront investments, though some ongoing retraining and fine-tuning costs remain, whereas inference costs scale continuously as usage picks up. Once a model is deployed to production, every single query or interaction incurs ongoing expenses, which accumulate fast as adoption continues to grow.

Also, at the beginning, a few major players dominated foundation-model training. Recently, however, a growing number of AI labs and open-source efforts are also making significant contributions and advancements–though these still account for a small percentage of all the companies using LLMs for business. The vast majority either fine-tune existing models or enhance them with techniques like Retrieval-Augmented Generation (RAG), which can significantly reduce the cost compared to training a foundation model from scratch.

Meanwhile, the demand for inference is nonstop. As GenAI applications proliferate and are applied to more and more industries, the computational burden of LLMs in production is skyrocketing. The rise of AI agents and complex workflows further amplifies this trend, as these systems often require multiple LLMs and embedding model calls within a single workflow. This compounding effect drives inference costs even higher, making efficient model serving a critical challenge for the future of AI.

By optimizing LLM serving, even using the same hardware and achieving similar latency, businesses can potentially achieve more throughput, serving more customers at the same time. This reduces hardware usage when they horizontally scale the serving system, saving a tremendous amount of money! Another example: an optimized model might be able to run on a lower-grade chip instead of a more advanced and expensive chip, while keeping similar latency and throughput. This also saves money.

Scalability, peak load handling, and feasibility

When a model is deployed in production, GPU demand scales with customer growth. An optimized inference solution enhances customer experience and reduces GPU costs, as you’ve seen. It also determines whether the system can scale efficiently under heavy traffic, especially when GPU availability is limited–and the industry is indeed just experiencing that.

For example, let’s say an LLM-powered sales agent runs smoothly, with stable traffic, most of the year. However, on Black Friday, demand can surge by 400% or more. A less optimized system may really struggle to scale to handle the sudden increase in traffic, leading to bottlenecks, degraded latency, and even request failures. The ability to handle peak traffic reliably is crucial for a resilient and scalable production system. This is a real-world problem we have personally experienced before, and getting the setup right can require a lot of tuning and performance testing to ensure service availability in those periods.

As mentioned in the last section, optimized models can be deployed on lower-grade chips, providing greater flexibility. Even with access to major cloud providers like AWS, Google Cloud, and Azure nowadays, high-end GPUs are not always readily available in every region. The ability to run models efficiently across a broader range of hardware, rather than being constrained to specific high-end GPUs, can be a key enabler for businesses expanding into new markets. We have frequently encountered this issue when deploying applications around the world.

The Role of Accelerator Chips in LLM Serving

Now that you understand the importance of LLM optimization, the next step is to explore the hardware that powers LLMs. Choosing the right accelerator setup is often one of the most important decisions in serving LLMs, because a system’s hardware constraints dictate its memory capacity, computational power, and efficiency. Understanding these factors will help you maximize your model’s performance and cost-effectiveness.

In this section, we’ll start by breaking down how to read GPU specifications. GPU specs can be overwhelming due to the sheer amount of technical details, but to simplify things, we’ll focus on the attributes most relevant to LLMs: GPU compute power, memory capacity, memory bandwidth, and interconnects. Understanding these factors will help you make informed decisions when selecting a GPU for your LLM inference and training.

After covering those key concepts, we’ll analyze and compare the specs of several popular NVIDIA GPUs (as of the time of writing, in early 2025). We’ll also share our insights on which GPUs are best suited for different LLM use cases.

We’ll end the section by looking at hardware options other than Nvidia chips, such as TPU and NPU.

Reading GPU specs

In this section, we will show you how to read GPU specifications by analyzing several attributes of two Nvidia GPUs: H100 SXM and H100 NVL. Understanding GPU specifications will not only help you choose the best hardware for AI inference, it will also benefit your model training.

Table 5-1. Specification comparison between H100 SXM and H100 NVL. (Source: NVIDIA)

H100 SXM
H100 NVL

FP64

34 teraFLOPS

30 teraFLOPS

FP64 Tensor Core

67 teraFLOPS

60 teraFLOPS

FP32

67 teraFLOPS

60 teraFLOPS

TF32 Tensor Core

989 teraFLOPS

835 teraFLOPS

BFLOAT16 Tensor Core

1979 teraFLOPS

1671 teraFLOPS

FP16 Tensor Core

1979 teraFLOPS

1671 teraFLOPS

FP8 Tensor Core

3958 teraFLOPS

3341 teraFLOPS

INT8 Tensor Core

3958 TOPS

3341 TOPS

GPU Memory

80 GB

94GB

GPU Memory Bandwidth

3.35 TB/s

3.9 TB/s

GPU Compute Attribute

GPU compute performance is usually measured in teraFLOPs. Floating-Point Operations Per Second (FLOPs) are a measure of a chip’s theoretical GPU compute capability, based on how many floating-point operations (like addition and multiplications) it can perform within a second. In other words, it’s a measure of how fast it is, theoretically. Tera- means 1 trillion, or 1,000,000,000,000. For example, the top seven attributes listed in Table 5-1 FP64, FP64 Tensor Core, FP32, FP32 Tensor Core, BFLOAT16 Tensor Core, FP8 Tensor Core, INT8 Tensor Core – represent the compute performance of different precisions in teraFLOPs.

From the comparison table, you can see that the H100 SXM can perform 1979 * 10¹² operations in a second, giving it FP16 precision! (FP16 stands for “half-precision floating point,” more often called 16-bit. You can double that if you can use a model with FP8 precision. FP16 and FP8 (8-bit) are lower-precision numerical formats used in deep learning models. Because FP8 is half the size of FP16, computation can be almost twice as fast. (We will discuss this a lot more in Chapter 6.)

Comparing the FLOPS horizontally between two versions, H100 SXM and H100 NVL, the NVL version is slightly less performant in this aspect. But is it really just worse for all use cases? Maybe not. Let’s continue our analysis.

GPU Memory Attributes

Next, let’s look at two equally important, if not more so, GPU specifications: memory (also called VRAM), which indicates the memory capacity for loading a model, and memory bandwidth, which measures the data-transfer speed for GPU computation. H100 NVL has a bigger memory size (94GB) than H100 SXM (80GB). What’s more, it also has a higher GPU memory bandwidth: 3.9 TB/s (terabytes per second), compared to 3.35 TB/s in SXM.

Now we have a choice to make: H100 SXM or H100 NVL? H100 SXM calculates faster than H100 NVL, since it has higher FLOPS, but H100 NVL has a better memory capacity and bandwidth. How do we choose between the two?

Let’s use an analogy here. Imagine the GPU is a big pizza kitchen, and our goal is to prepare pizzas as quickly as possible to serve as many customers as we can. Compute power, measured in FLOPS, represents how fast and powerful our oven is, while the GPU memory bandwidth determines how efficiently we can prepare all the ingredients and supply the dough to the oven. GPU memory capacity is the fridge’s capacity to store all the dough.

Do you see the interaction between all three? If we don’t have a big enough fridge to hold the dough, customers will now be sitting there hungry–just as, with a small GPU memory, we cannot load the model. But let’s suppose our fridge is just big enough to serve all of our customers. Let’s look at the oven.

If we have a very advanced oven, but we can only supply dough very slowly, our powerful oven will be running without baking much pizza. We can’t make good use of its power! We’ll have a similar problem if we have super high FLOPs for compute, but low GPU memory bandwidth.

Conversely, if the oven isn’t very powerful but we can prepare dough very fast, the pizza will take a very long time to bake and customers will have to wait for their food. This is much like having weak compute but high memory bandwidth.

Thus, it’s very important to find a balance between compute power, memory size, and bandwidth that fits your workload. We will discuss this more in the next few sections.

GPU Interconnect Attributes

GPU interconnects enable high-speed data transfer between multiple GPUs, both within a single node and across multiple nodes. It used to be crucial only for model training, but GPU interconnect has become more and more important for model serving.

With the emergence of larger and larger models (LLMs), one GPU is no longer powerful enough to host most models. You might need to serve a very large model which needs multiple GPUs, with each loading a portion of the model weights or layers. Or when dealing with a latency-sensitive use case, if one GPU is not powerful enough to achieve the latency requirement, you’ll need multiple GPUs to work together to achieve faster inference speed. When you face these situations, the GPU interconnect specs become critical. In this section, we will discuss GPU interconnect in detail, starting with intra-node interconnects (those within a single node) and then moving to inter-node interconnects (those across multiple nodes).

Intra-Node Interconnects

To help you understand intra-node interconnects, let’s look at three variations of the Nvidia GPU H100 in Table 5-2.

Table 5-2. Comparing GPU interconnects across three types of Nvidia GPU H100.

H100 PCIe
H100 NVL
H100 SXM

Form Factor

PCIe

PCIe

SXM

NVLink Support

No (Optional)

NVLink Bridge

NVLink/NVSwitch

Interconnect GPU-to-GPU bandwidth

128 GB/s with PCIe

600 GB/s with NVLink Bridge connecting 2 H100 GPUs only

900 GB/s connecting up-to 8 H100 GPUs

The first row of Table 5-2 tells you the form factor of the GPU. Form factor defines a GPU’s physical size, power requirements, and cooling design, which eventually determine how it integrates into a server or workstation. In this example, H100 SXM follows the SXM form factor. SXM-version GPUs are mounted directly onto the motherboard via a custom socket, rather than being inserted into a PCIe slot. This configuration enables faster connections, better power delivery, and enhanced cooling to achieve better performance with intense multi-GPU workloads.

In contrast, both H100 NVL and H100 PCIe follow the Peripheral Component Interconnect Express, or PCIe, form factor. PCIe is a commonly used standard data-transfer connection between various electronics. A PCIe version of a GPU means the GPU is installed into a PCIe slot, making it more widely compatible and cheaper, but less performant.

The next two rows show the NVLink technologies each variant uses and the corresponding intra-node GPU-to-GPU bandwidths. NVLink is a Nvidia high-speed interconnect technology that enables direct GPU-to-GPU interconnect for fast, high-bandwidth communication. For the variants without NVLink, GPU-to-GPU communication falls back to PCIe, which is much slower.

Let’s look at some GPU interconnect topology diagrams (shown in Figure 5-4 through Figure 5-7), to understand the difference between NVLink and PCIe, as well as different configurations of NVLink using these three versions of Nvidia H100 GPUs.

H100 PCIe

In Figure 5-4, two GPUs use PCIe, which has GPU-to-GPU bandwidth of 128 GB/s. (H100 PCIe can be installed with NVLink Bridge but needs to be bought separately.) This configuration is the cheapest. If you do not need fast GPU-to-GPU communication, such as if you have two small models that run independently on one GPU each, a basic PCIe connection is sufficient.

Figure 5-4. GPU interconnect with PCIe.

H100 NVL

H100 NVL also follows the PCIe form factor, but adds NVLink Bridge. This technology connects a pair of H100 NVL interconnects together with a bandwidth speed of 600GB/s, as shown in Figure 5-5.

The drawback of the NVLink Bridge setup is that it can only connect two GPUs. If you need more than two GPUs to work together, the slower PCIe connection is still required. For example, in a four-GPU setup, GPU1 and GPU2 are linked with NVLink Bridge, and GPU3 and GPU4 are also linked with NVLink Bridge. This means that communication between GPU1 and GPU2 takes place at 600 GB/s. The same goes for GPU3 and GPU4. But the connections between GPU1 and GPU3, and between GPU2 and GPU4, are still using PCIe, which is 128 GB/s. This setup is a good middle ground that balances cost and performance.

Figure 5-5. Two-GPU (top) and four-GPU (bottom) interconnects with NVLink Bridge.

H100 SXM

With the SXM version, up to eight GPUs within the same node can be connected with NVLink, as shown in Figure 5-6.

Figure 5-6. GPU interconnect with NVLink. (Source: NVIDIA Technical Blog)

Figure 5-6 shows how GPUs can be connected with NVLink in SXM. Since SXM can connect eight GPUs and each GPU’s interconnect bandwidth is 900 GB/s (see Table 5-2), each GPU must split its 900 GB/s total bandwidth into seven dedicated 128 GB/s (900GB/s divided by 7) point-to-point connections, each connecting to one of the other GPUs in the system.

If we don’t like the fact that we are splitting the interconnect bandwidth by the number of GPUs it’s connected to, we can use NVSwitches to achieve full 900 GB/s bandwidth with all-to-all GPU connections. NVSwitch is a separate, expensive, high-bandwidth switch that connects multiple NVLink-enabled GPUs together. Figure 5-7 shows that, with four NVSwitches added to the configuration in between, SXM can achieve full 900 GB/s bandwidth in all-to-all fashion for the best performance possible.

Figure 5-7. GPU interconnects with NVLink and four NVSwitches in between. (Source: NVIDIA Technical Blog)

Although using an NVSwitch increases interconnect bandwidth greatly, it is the most expensive of all the setups, and can be overkill if your workload does not require that much GPU-to-GPU communication.

Inter-Node Interconnect

All of the setups we just discussed are for intra-node GPU-to-GPU communication. However, it is not possible to have an unlimited number of GPUs in one node. The maximum number of GPUs within a node is usually capped at eight due to physical constraints, power supply, cooling concerns, and software support. As the LLMs become larger and larger, eight GPUs in one node can also become insufficient. In this case, you can shard models across multiple nodes, so even more GPUs can work together to serve the model.

Similar to intra-GPU interconnect, inter-node interconnect becomes important when you’re serving models across multiple nodes. One of the best solutions is using InfiniBand(IB) with GPUDirect RDMA. RDMA, which stands for “remote direct memory access,” enables direct memory access between GPUs across nodes. NDR 400G InfiniBand, for example, can achieve 50GB/s across nodes–but that’s still quite a lot slower than intra-node performance.

See Table 5-3 for a performance comparison of different intra-node vs. inter-node bandwidth. For example, within one node, GPU-to-GPU communication can be 900GB/s with NVLink and NVSwitch while GPU-to-GPU intra-node communication only maxes at 50 GB/s.

Table 5-3. Comparing communication speeds for example H100 setups.

Setup
Bandwidth

GPU-to-GPU within node with NVLink/NVSwitch

900 GB/s

GPU-to-GPU within node with NVLink Bridge

600 GB/s

GPU-to-GPU within node with PCIe

128 GB/s

GPU-to-GPU across node

50 GB/s

For LLM serving, the majority of deployments still run on a single node with one to eight GPUs for each model instance (also called a replica). As user traffic increases, we scale up the number of model replicas horizontally, by adding more hardware instances with the same setup. The goal of this approach is to minimize unnecessary internode communication.

In Chapter 6, we will explore how models can be distributed across multiple GPUs and multiple nodes and cover the pros and cons of different setups and techniques, such as Tensor Parallel and Pipeline Parallel.

Recent research is exploring separating the Prefill and Decode phases across different GPUs. Additionally, for large models using a Mixture-of-Experts (MoE) architecture—like DeepSeek V3/R1—techniques like Expert Parallel and Data Parallel are applied even across nodes. We will introduce these advanced topics after covering the foundational concepts.

Comparing the Specs of Popular GPUs

Now that you understand the key GPU specifications, let’s explore how to apply this knowledge to determine the best use cases for a given GPU.

Table 5-4. Comparing the specs and costs of popular GPUs. (These costs were current as of spring 2025 and may change.) (Source)

H200 SXM
H100 SXM
A100 SXM
L40S
A10

GPU Memory Size (GB)

141

80

80

48

24

FP16 and BF16 TeraFLOPs (TensorCore)

1979

1979

312

362

125

GPU Memory Bandwidth (TB/s)

4.8

3.35

1.935

0.864

0.6

Support for FP8

Yes

Yes

No

Yes

No

NVLink/NVSwitch

Yes

Yes

Yes

No

No

Cost on-demand hourly on AWS us-west1 or us-west2 (8 chips per instance)

P5e.48xlarge

Not available, only reserved

P5.48xlarge

~$123

P4de.24xlarge

$41 (80GB)

G6e.48xlarge

~$30

G5.48xlarge

~$16

Table 5-4 compares the specs and cost of some popular Nvidia chips for inference. In general, more powerful GPUs will always be more expensive, so your selection really depends on whether the model to be served can benefit from the GPU’s increased performance or the specific features it provides. Let’s look at some examples:

Small model

If you are serving a fine-tuned model such as Llama-3-8B for a specific use case and latency is not very critical, then A10 might be a decent choice. It already has enough GPU memory to load and execute the model, it’s the cheapest option here, and it’s highly available, being an older-generation chip.

Mid-sized model

The deepseek-ai/DeepSeek-R1-Distill-Qwen-14B model, which is a 14B model, has almost twice as many parameters as Llama-3-8B. If you’re serving this model, stepping up to L40S might be a good idea. Especially considering their support of FP8 precision, L40S can be highly capable for this task. his setup avoids all the overhead of serving models across multiple GPUs; at the same time, it has a good amount of memory left over to achieve long input context, long output generation, and high batch size.

Large model

If you need to serve deepseek-ai/DeepSeek-R1, NVLink Switch becomes a crucial feature to have. This model has 671B parameters, making it too big to fit and fast for even a single H200 to handle. A likely configuration would be using one machine with eight H200 GPUs connected together in FP8 precision. (The DeepSeek model is that it uses a Mixture-of-Experts (MoE) architecture. This enables further optimizations, such as having experts spread across multiple GPUs or even different nodes. We’ll discuss this in the final chapters of the book.)

These are just some simple examples of how different models run on different GPUs. We’ve introduced quite a lot of new concepts and calculations here, without expanding on them, but don’t worry. We will dissect them all in this book,and teach you techniques for finding the best solution that meets your serving requirement and reduces the cost-to-serve.

Please note that, in this hot field, technology changes fast. The latest, hottest chip that everyone wants today can be replaced by a newer one tomorrow, but the foundational concepts you’re learning and the intuition you’re building for weighing your options do not change.

Other AI Accelerators

So far, we have been solely investigating Nvidia GPUs. However, in the LLM Inferencing space, a lot of other chips are competing with Nvidia, and some are designed specifically for model inferencing. GPUs such as AMD’s MI300X and Intel’s Gaudi2 are both alternatives to H100 and A100. Tech giants Google’s TPU and Amazon’s Inferentia are mostly only available on their respective clouds. Huawei’s Ascend NPU is another notable competitor, particularly in markets where Nvidia faces restrictions. And there are many startups in this space, such as Groq, Cerabras, Untether AI, SambaNova, and d-Matrix.

As we write this in spring 2025, Nvidia’s General-Purpose Graphics Processing Unit (GPGPU) still dominates this market. There are several reasons for its dominance:

Some of these offerings are specifically tailored for deep-learning-based model inference, matrix multiplication workloads, or Transformers. Their architectures are simpler, more cost-efficient, and more energy-efficient, but compared to GPGPU, they often struggle with the rapidly evolving demand from model architectures.

At the software level, Nvidia GPUs have much more mature support with the CUDA ecosystem, which is widely used in both training and inference. Custom chips usually require their own proprietary software stacks, such as ROCm for AMD GPUs, JAX for TPU, and Neuron for Inferentia. Switching over can be costly, especially given the lower levels of community support,

Some chips that leverage on-chip SRAM to achieve impressive performance with superb latency. However, so far, this often isn’t a very cost-effective solution, given the high price of SRAM and the number of chips needed to serve one model instance.

Nvidia GPUs still have the edge in flexibility to meet different inference setups. They support different levels of precision, such as FP8, high memory bandwidth, and advanced GPU interconnect capability, which are all contributing factors in this area.

While all these factors contribute to Nvidia’s continued dominance in the market, the broader landscape of hardware acceleration is evolving rapidly. As new architectures emerge, one of the biggest challenges is overcoming bottlenecks, since data movement has not kept pace with improvements in raw compute power.

Figure 5-8 compares the increases in hardware compute FLOPS, memory bandwidth, and inter-GPU bandwidth over the last 20 years. You can see that the compute capability has improved a lot over time, and at a much faster pace than data movement speed, in both memory bandwidth and inter-GPU bandwidth. Nowadays, GPU data-movement limitations have become an obstacle to utilizing the fast advancements in compute. This problem is usually referred to as the “memory wall.” Quite a lot of accelerator chip manufacturers are trying to tackle this issue.

Figure 5-8. Scaling in peak hardware FLOPS and memory and interconnect bandwidth. (Source: An et al., 2024)

Now that you have a solid understanding of GPU specifications, it’s time to explore the limitations and bottlenecks you may encounter when actually serving LLMs on GPUs. Identifying these bottlenecks is critical, since the optimization techniques introduced throughout the rest of the book are specifically designed to overcome these challenges and improve efficiency.

Bottlenecks in LLM Model Loading

In the following two sections, we will put LLMs into the mix and see how they interact with the GPU specs you learned about in the two previous sections. We will start with model loading, then discuss model execution phases.

We’ll look at why models need to be hosted in GPU memory, identify the major consumers of GPU memory, and walk through how to estimate GPU memory requirements for a given model with a given workload.

The Model Loading Process

Figure 5-9. Data flow during model-weights loading.

The first thing you need to do for efficient LLM serving is to load the model weights and cache them in GPU memory. During the loading process, the model weights are first moved into CPU memory (also called system memory) and then transferred to GPU memory, as shown in Figure 5-9. Once loaded, the weights remain cached in GPU memory, ready to serve incoming requests. Thus, it’s essential to have a big enough GPU memory size to hold the model weights.

Some of you may wonder: what if we cached the model in a bigger CPU memory, or just load it on-demand when serving requests come in? The answer lies in data transfer speeds. Both hard disks and CPU memory are significantly slower at moving model weights, making them impractical for real-time inference. Table 5-5 compares the data-transfer bandwidths of SSD, CPU memory, and GPU memory. Loading from the hard disk would be even slower.

Table 5-5. Comparing data-transfer bandwidth between SSD, CPU memory, and GPU memory.

Hard Disk (SSD) bandwidth
CPU memory bandwidth
GPU memory bandwidth

0.5 to 14 GB/s

50 to 200 GB/s

300 GB/s to 3 TB/s

It’s crucial to distinguish between CPU (system) memory and GPU memory. The GPU memory is where the model weights should be cached. If model weights are stored inside CPU memory, they need to be transferred to GPU memory to utilize the GPU compute. This extra step introduces an unacceptable delay in serving responses.

Estimating Model Size

How much GPU memory will a large AI model require? When calculating the memory footprint of a model, there are two key variables: the model’s parameter count and the data type (precision) of its parameters.

Let’s take a model from HuggingFace as an example. A model’s number of parameters can often be found directly in its name. For instance, the name Llama-2-7b indicates that this model contains approximately 7 billion parameters in its model weight.

The data type can be found in the config.json file in the model’s Hugging Face repository (shown in Figure 5-10). Specifically, the torch_dtype attribute reveals the precision used for model weights. For example, if the value is float16 (FP16), this means the model uses 16-bit floating-point precision for storage and computation.

Here is an example of a config.json file in HuggingFace, which contains the model configuration and information to correctly instantiate and use the model:

{
  "_name_or_path": "meta-llama/Llama-2-7b-chat-hf",
  "architectures": [
    "LlamaForCausalLM"
  ],
  "bos_token_id": 1,
  "eos_token_id": 2,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 11008,
  "max_position_embeddings": 4096,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 32,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "tie_word_embeddings": false,
  "torch_dtype": "float16",
  "transformers_version": "4.32.0.dev0",
  "use_cache": true,
  "vocab_size": 32000
}

The data type, or precision, of a model refers to the number of bits required to store a single parameter and their corresponding memory sizes. Data type significantly impacts both model size and model performance, because it determines how the model weights are stored. In general, lowering precision—for example, from FP32 to FP16, and then to INT8 or FP8—comes at the cost of some accuracy, but reduces model size and improves serving performance. We will dive deeper in the next chapter; for now, you just need to understand the relationship between precision levels, as shown in Table 5-6.

Table 5-6. Parameter precision and memory size.

Bits
Size in bytes

FP32 (Single Precision)

32

4

FP16 (Half Precision) / BF16

16

2

INT8 (Quarter Precision ) / FP8

8

1

Now we know that the Llama-2-7b model has around 7 billion parameters that are stored in BF16 format, which means each parameter occupies 2 bytes:

~7 billion parameters * 2 bytes/parameter = 14 billion bytes = 14 GB

Looking at the files for the model confirms this estimate. The screenshot in Figure 5-10 shows the actual sizes of pytorch_model*.bin files, which sum up to around 13.GB.

Figure 5-10. Llama-2-7B model files in the HuggingFace model repository.

Estimating KV Cache Size

Continuing the example in the last section, we now know the model size, which is around 14GB. Suppose you have a GPU with 16GB of memory. You should be able to load it fine, since 16 is more than 14. It will very likely also run fine when you send in one short single request. Yay! But hang on. A GPU with just enough memory to hold the model is not ideal.

Figure 5-11. GPU memory usage breakdown.

Recall that in Chapter 2, we dived deep into KV caching, which is an optimization technique that trades some amount of GPU memory space for much faster serving performance by caching the intermediate results. As it decodes the next token, the model can reuse those results without redoing the computation. Figure 5-11 shows an example of how model weight and KV cache co-locate inside GPU memory. A small amount of extra memory should be set aside for intermediate computations. If you only leave a tiny amount of GPU memory for the KV cache, it constrains you to run at very limited batch sizes and with short context. Now let’s look at why, and see how to estimate how much KV cache space to reserve.

The formula for estimating KV cache size for each token is:

KV cache size per token = 2 x Number of Layers x Number of Attention Heads x Head Dimension x Data Type Size

Next, we need to account for the effects of the model’s architecture and size. In this simple example, we will use Llama-2-7b, which is still using the basic form of multi-head attention (MHA). In later chapters, we will introduce other novel ideas to reduce and compress KV caches from multi-query attention (MQA), grouped-query attention (GQA), and multi-head latent attention (MLA), which was introduced in DeepSeek V2 and is used in DeepSeek V3 and R1.

Llama-2-7b has 32 attention-head layers and 32 attention heads, with a head dimension of 128 (4096/32). Suppose we still use half precision. Referring to Table 5-6, you can see that the precision comes to2 bytes per token. Using the formula above, we get:

KV cache size per token= 2 x 32 x 32 x 128 x 2 = 524,288 bytes = 0.5 MB

We now know that each token requires 0.5MB. The next question is: how many tokens do we need to cache? The maximum number of tokens can be calculated as the maximum batch size times the maximum sequence length. Max batch size determines how many parallel users or /requests one model instance can process at a time. Max sequence length limits the total of input length and output length. Thus, if we are processing more requests together and dealing with longer prompts and/or longer generation, the KV cache size will add up:

Total KV cache = KV cache size per token x max number of tokens in cache

= KV cache size per token x (max batch size x max sequence length)

Suppose we use the model to summarize long paragraphs, using a max sequence length of 4,096. With a batch size of 16, the calculation would be:

Total KV cache = 0.5 MB/token * 4,096 tokens per request * 16 requests = 32GB

That’s even bigger than the model itself, which is now 14GB!

Now let’s consider two chips, with 24GB and 48GB GPU memory, respectively, still with the max sequence length of 4,096. As shown in Table 5-7, using total GPU memory size and subtracting the model size, which is around 14GB, we can calculate the estimated memory left. We can then derive the batch size by dividing the product of the sequence length (which is 4,096) and 0.5MB per token. In this case, theA10 GPU can only serve 4 parallel requests, while L40S can serve 16. Even though L40S is more expensive, it still comes out to be more cost-efficient. (Note, however, that we usually cannot achieve the maximum theoretical batch size from the batch size calculation, because we also need to reserve space inside the GPU for activations, which are the tensors created in the intermediate steps.)

Table 5-7. Batch size calculation and cost comparison between two different GPUs.

GPU and Memory Size
Memory Left After Model Weight Loaded
Max Possible Batch Size
On-demand hourly cost per chip from AWS

A10 24GB

10GB (24 - 14)

4 (10*1024/(0.5*4096) = 5 )

$2

L40S 48GB

34GB (48 - 14)

16 (34*1024/(0.5*4096) = 17)

$3.75

So far, you’ve learned that during model loading, we want to reserve enough space for both the model weights and KV cache. Model weights need to be cached inside the GPU memory; KV cache space is reserved for execution. Figure 5-12 shows an example of how GPU memory usage changes from idling to execution. During execution, KV cache size will continue to grow as the sequence length grows and we cache more tokens. What we need to make sure is we have enough memory beyond the peak memory usage at the very end of the generation to avoid Out-of-Memory (OOM) errors.

Figure 5-12. How GPU memory usage changes from idling to execution.

In later chapters, we will show you techniques (such as prefix caching) that may require more space to achieve faster inference, especially a lightning-fast time to first token (TTFT). But as a general rule of thumb, when estimating model GPU memory requirements, we recommend getting started with GPU memory approximately twice the model size to achieve better parallelism and unleash more GPU performance.

Bottlenecks in LLM Model Execution

LLM execution is a complex process. Recall what you learned in Chapter 2: when an LLM processes input and then generates output, it goes through two distinct phases: Prefill and Decode. These phases exhibit significantly different GPU utilization patterns. To analyze this, we will introduce the concepts of arithmetic intensity, which measures computational efficiency, and the roofline model, which illustrates hardware limitations and performance constraints.

Boundaries of GPU Compute and Memory Bandwidth

With model loading and its relationship with model size out of the way, now we have the model entirely loaded inside GPU memory and ready for model serving. It is time to discuss the elephant in the room: is the model serving bounded by GPU compute FLOPs or GPU memory bandwidth? In other words, recalling the pizza kitchen analogy, do we need a bigger oven that can bake faster, or do we need to prepare the ingredients and dough faster? Which of the two is the bottleneck in our model serving workload?

To analyze that, we need the concept of arithmetic intensity: the ratio of algorithm implementation operations to the number of bytes accessed. This basically calculates the compute (FLOPS) over data movement to the compute unit in a ratio of FLOPS per byte.

Arithmetic Intensity = Number of FLOPS / Data Movement

If a workload requires very low computation but needs a lot of data reading and writing, then it has a low arithmetic intensity. On the contrary, if a workload does not read and write a lot of data but performs a lot of computations on a small amount of data, it has a high arithmetic intensity.

Now, let’s look at data movement, which is often confused with model loading. Data movement happens at model execution time, when the model weights are already in GPU memory. During serving, the model weights are moved from GPU memory all the way to the register, so the actual compute can execute the operation (Figure 5-13).

Figure 5-13. Data movement during model serving, as model weights are loaded to register for computation.

During model serving, the model weights stored inside GPU memory are constantly read, along with all intermediate step outputs to registers for computation. Since the GPU memory is the slowest of the bunch, GPU memory bandwidth is the number we should use for the data movement to the compute unit. (It’s a little different if the weights or outputs are small enough to fit into the on-chip caches in a smarter way, such as flash attention, but we’ll talk about that in Chapter 6.)

Just by looking at the GPU specs, such as the L40S GPU specs in Table 5-8, we can calculate the chip’s theoretical arithmetic intensity.

Table 5-8. Specifications for the Nvidia L40S GPU.

GPU Memory
48GB GDDR6 with ECC

Memory Bandwidth

864 GB/s

Interconnect Interface

PCIe Gen4 x16: 64 GB/s bidirectional

NVIDIA Ada Lovelace Architecture CUDA Cores

18,176

NVIDIA Third-Generation RT Cores

142

NVIDIA Fourth-Generation Tensor Cores

568

RT Core Performance TFLOPS

212

FP32 TFLOPS

91.6

TF32 Tensor Core TFLOPS

183

BFLOAT16 Tensor Core TFLOPS

362.05

FP16 Tensor Core

362.05

FP8 Tensor Core

733

Peak INT8 Tensor TOPS

733

For half precision (FP16 Tensor Core), we can calculate the theoretical arithmetic intensity for the chip:

362 TeraFLOPs / 864GB/s = (362 * 10^12 FLOPs) / (864 * 10^9 Bytes) ~= 419 FLOPS/B

Now that we have the arithmetic intensity, we can draw a roofline model, a visual performance model that relates a system’s compute capability and memory bandwidth to help identify whether an application is compute-bound or memory-bound. A naive roofline model for the L40S chip, a simplified version that estimates a system’s achievable performance based only on its peak compute throughput and memory bandwidth, assuming ideal conditions without accounting for hardware complexities, is shown in Figure 5-14.

Figure 5-14. Roofline model for the L40s chip.

The x-axis shows the arithmetic intensity in FLOPS/B while the y-axis shows teraFLOPs (TFLOPS). Figure 5-14 shows how much compute the GPU can theoretically do given a certain arithmetic intensity in FLOPS/B. If the arithmetic intensity is lower than the 419 FLOPS/B we calculated, then we cannot use all the 362 TFLOPS the GPU is capable of.

It’s time to try using arithmetic intensity and a roofline model to analyze whether a specific workload on a given GPU is bounded by compute or memory bandwidth. Let’s say we have a workload to run with ~210 FLOPS/B, basically half of the crossover point of 419 FLOPS/B. This is depicted in Figure 5-15, data point 1. In data point 2, where the workload has 1,000 FLOPS/B, which easily hits the maximum TFLOPS the chip can do. In this case, data point 1 is bound by memory bandwidth, because it cannot utilize the extra compute it has available. (We still have a lot of space in the pizza oven, but we can’t supply the dough fast enough.) Data point 2 is considered to be compute-bound because we’ve already reached the maximum compute available. No matter how much more data we can read, this is the fastest we can calculate using this chip. (We’re supplying dough very fast, but the oven has limited capacity and we have to wait for its current load to finish baking.)

Figure 5-15. Analyzing a workload and how it is bounded using arithmetic intensity and a roofline model.

Now, with a good understanding of the hardware limitations, we need to figure out what an LLM serving workload looks like. Is it memory bandwidth-bound with low arithmetic intensity, or compute bound with high arithmetic intensity?

Arithmetic Intensity in Matrix Multiplications

To understand whether the LLM serving workload is compute-bound or memory bandwidth -bound, we need to calculate the arithmetic intensity of the various layers inside the LLM’s architecture.

In Chapter 2, you learned that an LLM is mostly made up of Transformer blocks. Inside the Transformer blocks, there are two main components: self-attention layers and feedforward layers. The vast majority of computations for both are matrix multiplications. In LLM architectures, there are many other layers, such as element-wise operations and reduction operations, all of which typically have low arithmetic intensity because the computation on each data point loaded is not as high as in matrix multiplications. But these layers only represent a small portion of the overall computation. So, in this section, we’ll first show you how to estimate the arithmetic intensity of a matrix multiplication operation (often shortened to matmul).

Matrix multiplication means calculating dot products between rows of the first matrix and columns of the second matrix to produce a new matrix. Figure 5-16 shows an input matrix with shape [M, K] doing a matmul with a weight matrix [K, N], yielding the output matrix [M, N] shown on the right.

Figure 5-16. Generic matrix multiplication of input and weight matrices, producing an output matrix.

The naive code to calculate the matmul is shown below in nested triple for loops:

for m in [0, M):
	for n in [0, N):
		for k in [0, K):
			Outputs[m][n]+=Inputs[m][k]*Weights[k][n]

To calculate the arithmetic intensity, we need to calculate two things: the numerator, which is the number of operations performed during the matmul, and the denominator, which is how much data is transferred.

Numerator

For number of operations, based on the code snippet above, there are M * N * K multiplications and M * N * (K - 1) additions. We can simplify it to

# of Operations = 2 * M * N * K

Denominator

For data movement, both input matrices are needed to be read and output matrix is needed to be written. Assuming the precision of the values inside the matrices are all half precision, which has size of 2 bytes each. We can get

Data Movement = 2 * (Inputs size + Weight size + Outputs size)

= 2 * (M * K + K * N + M * N)

Thus:

Arithmetic Intensity = # of Operations / Data Movement

= 2 * M * N * K / (2 * (M * K + K * N + M * N))

= M * N * K / M * K + K * N + M * N

Using our formula for calculating the arithmetic intensity of a matmul, let’s see how the arithmetic intensity changes given different sizes of matrices. Supposing for simplicity that M = N = K , let’s calculate a few values for arithmetic intensity, as shown in Table 5-9.

Table 5-9. Calculating arithmetic intensity values for matrices of different sizes.

Matrix Size (M=N=K)

Arithmetic Intensity (FLOPs/Byte)

Verdict based on L40S

64

21

Memory Bandwidth Bound

512

170

Memory Bandwidth Bound

4096

1365

Compute Bound

The operation can be pushed to compute bound if the matrices in the matrix multiplication operation are large enough. Then the question in LLM serving is: Are the matrices large enough? You might think, given the size of the model, it should be, right? Not always. Let’s break that down in the next section.

Applying arithmetic intensity analysis to the LLM Prefill and Decode phases

You now know that, the bigger the compute matrix, the higher the computational intensity is during matrix multiplication. However, it’s not always possible to achieve big matrix multiplications in LLM serving. This is especially true when the request batch size is 1, which means the model is only processing a single request at the time. (We will discuss batching in Chapter 6.)

Let’s dive deeper into the shape of the input data passed to the model. The “shape” here describes the dimensions of the data, commonly referred to as its tensor. In this case, our input tensor is three-dimensional and has a shape of [batch size, sequence length, model hidden dimension]. Batch size is the number of requests we are processing at once. Sequence length is the number of tokens in the input prompt. From simple phrases such as “Hello” to big paragraphs a million tokens long, the sequence length is highly variable and mostly up to the user. Model hidden dimension is the number of values each token’s vector has in the model, impacting its capacity and performance. Since we are now only considering one request, batch size is set to 1. Consequently, the three-dimensional input tensor can be reduced to a matrix that has only two dimensions [sequence length, model hidden dimension].

Before we start the arithmetic intensity analysis, let’s revisit the two main phases of model serving from Chapter 2: Prefill and Decode. The Prefill phase is the initial processing stage, where the model encodes the input prompt before generating any output tokens. During this phase, the model does two things. First, it computes attention over the input tokens using its transformer layers; second, it generates the first set of KV cache entries and stores them for the subsequent generation phase.

In the next stage, the Decode phase, the model generates tokens one by one using previously computed KV cache.

Now let’s analyze the arithmetic intensity of the Prefill and Decode phases by looking at their matrix multiplication. You will see the difference in Figure 5-17.

Figure 5-17. Input size difference between the Prefill and Decode stages in a matmul operation.

In Figure 5-17, M, N, and K–which we used in Figure 5-16 for generic matrix multiplication–are replaced with h and s to demonstrate the actual shape of the matrices in LLM serving. The variable h corresponds to the weight (or “hidden”) dimension, which does not change between Prefill and Decode. In Prefill, because all tokens are processed together, the matrix has a large dimension for the number of rows in s, which is the sequence length of the input. During Decode, however, because tokens are generated one token at a time, s_decode, which denotes the sequence length during decode, is only 1.

Next, in Table 5-10, let’s plug the numbers back into the arithmetic intensity formula, which again is:

Arithmetic Intensity = M * N * K / M * K + K * N + M * N

= s * h * h / s * h + h * h + h * h

Table 5-10. Arithmetic intensity comparison between Prefill and Decode at different sequence lengths.

Sequence Length
M = s for Prefill
M = 1 for Decode
Hidden Dim
K = N = h
Prefill
Arithmetic Intensity (FLOPs/Byte)
Decode Arithmetic Intensity (FLOPs/Byte)
Verdict based on L40S

64

4096

31

0.5

Memory bandwidth-bound for both Prefill and Decode

512

4096

240

0.5

Memory bandwidth-bound for both Prefill and Decode

4096

4096

1365

0.5

Compute-bound on Prefill but Memory bandwidth-bound on Decode

From the numbers we get from for Prefill and Decode calculations, you can see that Prefill can have a high enough arithmetic intensity to saturate GPU compute, making Prefill compute bound when the sequence length is high. However, no matter how long the sequence is, Decode will always be very low, making it memory bandwidth bound.

In summary, different components such as Prefill vs Decode, and different layers of the model can all have very different arithmetic intensity. Arithmetic intensity is also affected by sequence length (as well as batch size, which we’ve intentionally skipped for now). All these analyses and calculations are designed to help you develop your intuition about bottlenecks in the different phases of LLM serving, so you’ll later be able to understand why various optimization techniques work and how to select them. For example, to improve performance, if a workload is compute-bound, you should explore ways to optimize mathematical computations and reduce FLOPs. If it is memory bandwidth-bound, you need to minimize unnecessary data movement.

Summary

This chapter began by showing you why LLM serving is crucial for business applications. Efficient LLM serving can improve customer experience, save costs, and even determine whether the business can be sustained and expanded.

Next, we turned to understanding AI accelerators, focusing on the key GPU specs of memory size, compute FLOPS, and memory bandwidth. You also learned why intra-node and inter-node GPU communication technology can be important for large models when a single GPU is not enough. Additionally, we touched on the current state of other AI accelerators and how they are trying to overcome current limitations on GPU advancements, known as the “memory wall.”

With that understanding of the hardware side, we started to show you how an LLM runs on the hardware in both the model loading and model serving phases. We introduced arithmetic intensity as a key approximation technique in estimating model performance and identifying if it is bottlenecked by compute or memory bandwidth.

After some concrete examples with simplified calculations to illustrate the change in arithmetic intensity during Prefill and Decode, we concluded that LLM serving has both compute-bound phases and memory bandwidth-bound phases. When batch size is fixed at 1, LLM can be very memory bandwidth-bound, especially when Decode or the input sequence are not long enough. In this case, the GPU compute is not saturated enough to reach its full potential.

In this fast-evolving field, new techniques come out at an unprecedented pace. But the intuition you built from analysis will equip you to better understand them and adopt them if they fit your situation. Even if language models switch to a new architecture or you start working on vision models, your knowledge about hardware specs, how models run on hardware, and the general intuition will remain valuable.

Chapter 6. Essential LLM Optimization Techniques

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 6th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

In the prior chapters, we have demonstrated the importance and the challenges of optimizing LLMs for serving. In the next two chapters, we will dive deep into each of the critical LLM optimization techniques one by one so that you are equipped with the knowledge to decide when, how, and why to use them for your serving needs.

In this chapter specifically, we will focus on essential techniques that will help you understand most optimization concepts and achieve a lot of your optimization goals. We’ll leave the more advanced techniques and industry trends for Chapter 7.

In this chapter, we will discuss how to use:

Request batching and scheduling to achieve better parallelism and GPU utilization

Attention optimization to achieve better compute efficiency, less required compute and better memory management

Model compression to achieve smaller model, less memory movement, and/or less compute

Prefix caching to cache and reuse prior prompts, including how to do it efficiently and obtain a high cache-hit rate

Request Batching and Scheduling-Level Optimizations

In Chapter 2, we divided serving into offline serving and real-time online serving. In real-time online serving, requests are received as the user sends them, while for offline serving, we have the requests already and can batch them all together so that they form a big tensor input that gets fed into the model, instead of sending them one by one.

Grouping requests together during serving can achieve higher throughput, even though it can mean slower responses. Let’s dive a bit deeper here to understand why, using a concept you learned in Chapter 5: arithmetic intensity.

Why do we need batching in real-time serving?

Recall that LLM serving has two phases: Prefill and Decode. The Prefill phase is when the model understands the input prompts. Tokens in the input prompt can be processed in parallel and achieve high arithmetic intensity, leading to a compute-bound workload. The Decode phase is when the LLM generates one token at a time. Because of this phase’s autoregressive nature, Decode is memory-bandwidth-bound: the model basically needs to go through all of its billion level parameters to generate only one token at a time, which is very inefficient in terms of GPU memory bandwidth (measured in FLOPs) (Figure 6-1).

Figure 6-1. Tokens in the input prompt are processed in parallel, all together in one go, during Prefill, while new tokens are generated one by one during Decode.

To fully utilize the GPU compute, you can add more requests with batching and process them together. Let’s add that into the setup shown in Figure 6-1. We’ll suppose a max batch size of 3 for all the figures in this section.

We’ll take three input prompts, prompt1, prompt2, and prompt3, batch them together, and send them to the model. The Decode step still only generates one token per iteration, but since we are batching these requests together, we can generate three new tokens in one go: one for each request, as shown in Figure 6-2. This artificially increases the arithmetic intensity: we are still reading the model weights once, but doing more calculations and generating more tokens.

Figure 6-2. With three batched requests, during Decode, three new tokens can be generated all together.

In summary, batching is especially effective during the Decode phase, where the model generates one token at a time. By grouping multiple requests, batching increases overall throughput and improves GPU FLOPs utilization. In contrast, its benefits are limited during the Prefill phase, during which the model already processes all input tokens in parallel. As long as the input prompt is not very small with very limited tokens, (say, less than 1,024), Prefill easily saturates GPU compute capacity on its own, so additional parallelism from batching has minimal impact.

Dynamic Batching in Online Inference

The example we just gave you showed the benefits of batching, but in real life for online inference, you won’t have a list of prompts you can process beforehand. So how do you achieve batching? One way is through client-side batching, where the client tries to accumulate requests on its end and send them all together once it reaches a preset batch size. We can do something similar on the server side: having the server wait for incoming requests to fill the preset batch size completely before processing them. This is called static batching.

Both methods are good for offline use cases, but not ideal in online inference. While offline batching is very common and straightforward, online batching is much more complex. This is because real-world requests are random and can arrive with big time gaps between them, which means filling a batch to maximum can take a long time, resulting in high latency. For example, suppose we have a batch size of 10. The first 9 of those 10 requests arrive within a single second, but the 10th request arrives 5 minutes later. The first 9 requests would not be processed until the 10th arrives; they would just wait there for 5 minutes.

The solution is dynamic batching, which groups incoming requests together at inference time based on defined logic, to balance latency and throughput dynamically. This logic usually uses two key parameters: batch size (also called max batch size, preferred batch size, or max number of sequences), and max delay time. Batch size, as you know, determines how many requests can be grouped together before they are sent to the model. Max delay time is the maximum time that existing requests can be held while waiting for the other requests to fill the max batch size. If the pending request count reaches the max batch size, the batch is sent immediately, even if the max delay time is not yet met. Similarly, if the max delay time is reached, the whole batch is sent immediately–even if there is only one request in it.

We can make an analogy here. Imagine you’re managing a small ferry-boat business. Each boat can carry up to 10 people across the river at a time. People arrive at the dock at different times, and your job is to minimize their waiting time while also trying to avoid sending half-empty boats. Without batching, you’d be sending a boat across every time someone arrives: basically one person, one boat. The passengers would have a great experience, but it would be a horribly inefficient and expensive way to run a business. With static batching, the boat would wait at the dock until 10 people showed up to fill it completely. This would be efficient, with minimal waste, but if people arrived slowly, the very first arrival might have to wait for a long time. That kind of latency would be just as bad for a ferry-boat business as it is for LLM serving. Dynamic batching tries to find a good middle ground that balances latency and throughput: the boat’s max batch size would still be 10, but you’d also set a max delay time of 5 minutes. If only 8 people arrive within that 5 minutes, the boat can leave without further waiting. If 10 people arrive before the 5-minute mark, the boat is sent right away.

Now that you know the benefit of dynamic batching, you also need to understand how to tune these two parameters–max batch size and max delay time–based on your latency and throughput requirements. Generally speaking, you want to achieve as high a batch size as possible while still meeting the latency SLA. High batch sizes aren’t always possible, because as you increase batch size, the processing latency will go up, and the memory usage of the GPU or CPU will also increase until it runs out of memory. Max delay time is another one to tune. Too long a delay time combined with high batch size would force requests that have already arrived to wait for an extended period of time; too short a delay time would prevent you from filling the batch in time and reduce the actual batch size sent for processing.

Continuous Batching for LLM Online Inference

Dynamic batching works pretty well for most models and has been adopted for many traditional ML serving settings. However, LLMs pose a bigger and unique challenge: LLM serving can have varying input and output lengths, which causes requests in a batch to take vastly different amounts of time to process. In dynamic batching, the total time it takes for the batch to finish is determined by the longest and slowest request, since the result is returned when all requests in the batch complete. Figure 6-3 is a simplified example where we abstract out prefill and decode, just to show you that because the requests have different lengths, one very long request can cause a lot of idle time.

Figure 6-3. The uneven lengths of the three requests can cause significant GPU idle time in LLM serving, as requests 1 and 2 wait for request 3 to finish

In the boat-business analogy, let’s assume your boat needs to take people directly to their homes on the other side of the river. Their trip times could be very different depending on where they live, just like LLM requests having different lengths. The boat has to drop off the person whose home is the furthest away before returning to the dock for the next batch of people, which means it will run with lower and lower capacity.

In order to mitigate this unique challenge, continuous batching (also known as inflight batching or iterative batching) for LLM serving was introduced. While dynamic batching has a wait time and processes requests batch by batch, continuous batching adds the requests to the backend model and group them on the fly, without waiting for a fixed batch size or time window. In other words, as soon as one running request in the batch finishes, the queued request gets added to the batch immediately.

In Figure 6-4, you can see that at the very beginning, requests 1, 2, and 3 are sent to the model for processing. Once request 1 finishes, a continuous batching strategy immediately adds the newly arrived request 4 into the processing. The same goes for the newly arrived request 5 when request 2 finishes, as well as request 6 when request 5 finishes. With dynamic batching, newly arrived requests 4, 5, 6 would all be waiting for request 3 to finish before they are dispatched for processing.

Figure 6-4. With continuous batching, once one request finishes, another is added for processing to minimize GPU idle time.

In the boat analogy, you have one boat that can hold a max of 10 people, and you must carefully manage this one boat to make sure it does not wait too long or run while empty. Now, what if you replaced that single 10-person boat with 10 small boats that can each hold one person? This way, once a person arrives, one boat could leave immediately. Even if the people have different homes at different distances, there is no waste: the boat just comes back once it transports that person, and picks up the next person waiting. This sounds like the perfect scenario!

With continuous batching, the next request in the queue is sent for processing once one prior request is complete. This means you don’t need to set a max delay time artificially, like in dynamic batching. You’ll need to manage and tune the max batch size, though, and a higher max batch size can still cause higher latency for individual requests, even as it increases overall throughput.

In many modern LLM serving frameworks, another parameter is added to control whether a request is queued or added to the batch: the max number of batched tokens. While max batch size controls things at the request level, basically determining how many requests the model is processing at a given time, capping the number of batched tokens provides more granular control at the token level. That’s important because input lengths for LLM serving can vary a lot. For example, batching 10 requests with an input length of just 20 tokens each would be a very different workload than batching 2 requests with 100,000 tokens each. Solely relying on the request level limit is not enough to take different token lengths into account and can lead to batching together too little or too much workload.

Figure 6-5. Demonstration of the relationship between max batch size, max number of tokens, and max model length

To better illustrate the relationship between these parameters, see Figure 6-5. Here, the max batch size is set to 3, so in total, 3 requests can be processed all together in a batch. The length of the bar for each request represents the number of tokens, which should be less than the maximum model length, also known as max context size, of the model (discussed in Chapter 5). The maximum number of tokens controls how many tokens in total the scheduler allows to be batched. If the scheduler receives a few very long requests that easily hit the max number of tokens limit, it will stop more requests from joining the batch.

With these two parameters, max batch size and max number of tokens, working together and requiring us to meet both sets of constraints, we can fine-tune how many requests can be processed across the Prefill and Decode phases. In the Prefill phase, the key constraint is the maximum token count, since inputs are much longer. In the Decode phase, parallelism is typically capped by the maximum batch size. It’s important to set the max number of tokens high enough. If it is too low, we might not saturate our GPU compute, because we won’t be sending enough tokens for the GPU to process in parallel during Prefill. Example 6-1 provides code to configure these parameters in vLLM.

Example 6-1. Code snippet to configure max number of batched tokens (--max-num-batched-tokens) and max batch size (--max-num-seqs) in vLLM.

vllm serve \
 Qwen/Qwen2.5-7B-Instruct \
 --max-num-batched-tokens 4096
 --max-num-seqs 128

Continuous Batching with Chunked Prefill

Continuous batching is great for solving the problem of requests having different lengths, but we are overlooking another unique aspect of LLM serving: the Prefill and Decode phases constitute two different workloads. In the last section, when we discuss dynamic batching and continuous batching, we abstracted away the Prefill and Decode phases in the diagrams and solely used bars to represent request lengths. Now it is time to take a deeper look at how continuous batching works and see if there are opportunities to improve further.

Recall that in LLM serving, Prefill is the phase where the model first processes the input prompt. This phase has high arithmetic intensity and does not require much help from batching to be efficient to run on GPUs. Decode is the phase where the model starts generating new tokens in an autoregressive manner; it usually has low arithmetic intensity and can benefit a lot from batching.

When we discussed the benefit of batching, we used a diagram similar to the one in Figure 6-6, which is a cherry-picked “happy path” scenario in which everything lines up perfectly: request input prompts all have the exact same length, generate the same number of tokens, and start running at the same time. In the first iteration, the Prefill steps from all three requests are batched together. Upon finishing that batched iteration, the model starts the second batched iteration, which contains the Decode steps from each of the three requests running together.

Figure 6-6. A cherry-picked “happy path” scenario where all requests have Prefill and Decode lined up perfectly

However, in the real world, not only are requests’ input and output lengths random, but with continuous batching, the start times are also different. Requests arrive at different times and we process them as soon as we can. So here’s the question: If the first request is already running in the Decode phase and the second request arrives and wants to start Prefill, do we prioritize the Decode of one request or the Prefill of another request? Can we batch them together in the same iteration?

Let’s start with the solution of not batching Prefill and Decode together, since running a hybrid workload often requires more complex GPU kernels. In this way, as Figure 6-7 shows, we only receive one request (request 1). We performed two iterations: iteration 1 to finish the Prefill phase, and iteration 2 to perform one Decode step. Then requests 2 and 3 come in. Now we need to decide whether to prioritize Prefill or Decode. Usually, we choose Prefill, because Prefill determines the time-to-first-token (TTFT), an important latency metric–especially for interactive usages such as chatbots. However, when we’re doing Prefill for requests 2 and 3 in iteration 3, request 1 will be waiting there completely idle. If the prompts for requests 2 and 3 are long, like those shown in iteration 3, it really hurts the end-to-end latency and inter-token latency of request 1.

Figure 6-7. Continuous batching with no Prefill and Decode, batched together and prioritizing Prefill

Now, let’s consider combining Prefill and Decode together, as shown in Figure 6-8. Again, the second Decode step of request 1 is moved into iteration 3, to run with the Prefill steps of request 1 and 2 in the same batched iteration. However, this still doesn’t help much. The delay is still quite prominent, because decoding one token can go much faster than finishing the Prefill step, especially when the input prompt is long.

Figure 6-8. Continuous batching, with Decode batched with Prefill together in iteration 3

One solution to mitigate this is called chunked prefill, a technique that splits the long input prompt into smaller chunks. As shown in Figure 6-9, all the prior long Prefill bars are now split into several smaller Prefill boxes with similar sizes to the Decode boxes (ideally), meaning they have similar processing times. As requests 2 and 3 join the batched iteration, request 1 continues to do the Decode step, while the other two requests start their own smaller chunked Prefill step, as shown in iteration 5. Later, in iteration 10, request 2 seamlessly transitions to Decode, since it has a shorter Prefill than request 3.

Figure 6-9. Continuous batching with chunked prefill

The benefit of chunked prefill is that even when the input prompt is long, because it is chunked into smaller ones, Decode boxes are no longer waiting there, blocked by the long Prefill. This improves inter-token latency (ITL), which we discussed in Chapter 4. However, it gains ITL at the cost of Prefill performance, since it adds more workload during Prefill, which makes the TTFT longer. Chunked prefill doesn’t really improve end-to-end latency either; in fact, it usually has some negative effects due to the overhead of computing multiple smaller Prefill steps. So this is a tradeoff you need to make based on your use case and SLA requirements. On the other hand, throughput usually improves with chunked prefill, because it improves batch efficiency by filling gaps of idle time to better utilize the GPU.

Another parameter you need to tune is how many smaller Prefills you split the original long Prefill into–in other words, how many tokens you want your small Prefill jobs to handle. This is also part of the max number of batched tokens. If you choose an extremely large number of tokens–for example, taking it all the way to the max model length would mean not doing any chunking of the Prefill at all. Too few tokens in each small Prefill would mean more overhead, which is not good either: we wouldn’t be batching enough tokens in one iteration to saturate GPU compute. Ideally, we would find a value in the middle: not small enough to cause a lot of overhead and low GPU utilization, and not large enough to defeat the purpose of doing chunked prefill.

Continuous batching has been the industry standard in production-level LLM serving for a few years as of this writing. Chunked prefill and its variations are also very popular for many use cases, such as handling long context workloads. One even more advanced technique, called Prefill-Decode disaggregation, completely breaks down the Prefill and Decode jobs into different GPUs and even different nodes. We will discuss this in the next chapter, after building more basic knowledge in this chapter first.

Scaling Attention and Kernel Optimization

Recall from Chapter 2 that during LLM serving, the two major components in a transformer block are the attention and feed forward layers. In this section, we focus on how to scale the attention aspect of LLM serving. We will start with reducing KV cache size, then discuss how custom GPU kernels improve attention computation and memory access.

The attention mechanism is the key to the breakthrough of LLMs. Yet the original multihead attention formulation that powers GPT‑3 is no longer the only or the most efficient way to leverage queries, keys, and values at scale. Production workloads and the high compute cost of LLMs force us to squeeze more performance out of models, spawning the creation of multi‑query attention (MQA), grouped‑query attention (GQA), and the latest, multihead latent attention (MLA), to slash KV‑cache size while preserving model quality.

After we look at attention optimization, we will discuss the GPU kernels that execute attention. Over time, these kernels have advanced from generic kernel fusion to attention‑specialized, hardware‑tuned, cache‑aware kernels. We will start with understanding kernel fusion and then go into specific kernels, such as FlashAttention.

At the end of the chapter, we’ll look at another key innovation, PagedAttention, which introduces a novel mechanism to manage how KV cache is saved in GPU memory to achieve high memory utilization.

Scalable attention mechanisms

Reducing the KV cache size is one key way to improve LLM serving performance. With smaller KV cache size, during serving–and especially during a GPU memory-bandwidth-bounded Decode phase–there will be less data movement. This alleviates pressure on GPU memory bandwidth, as you learned in Chapter 5. A smaller KV cache also means less space is used in GPU memory, which means you can run more requests in parallel with higher batch sizes for more throughput, and serve longer context than was previously possible with limited GPU memory space.

Figure 6-10. Comparing multihead attention, multi-query attention, grouped-query attention, and multihead latent attention

Starting with the very left of Figure 6-10, multihead attention (MHA) is the original version that underpins many earlier models. But in this original setup, for each query, calculating its attention requires one distinct key and value head. This makes the KV cache for that model architecture the largest and the least efficient of all four of the setups we will discuss here.

To improve performance, the third setup in the diagram, multiquery attention (MQA), was introduced. In MQA, all queries share a single key and value head. For a 7B model and 70B model, there are usually 32 and 64 attention heads, respectively. In MHA, this means it will have 32 and 64 KVs ,whereas in MQA, it will have only one. Thus, with MQA, KV cache size can be reduced by 32 and 64 times, which is a huge gain. However, MQA has been shown to degrade model accuracy significantly because of its aggressive setup.

To mitigate the accuracy issue of MQA, grouped-query attention (GQA) was introduced to balance performance and accuracy. GQA groups query heads into multiple groups, each with the same key and value. This setup has proven to be a good middle ground between MHA, which is not computationally efficient, and MQA, which takes too big an accuracy hit, and is now used in many model architectures.

Now let’s quickly examine how to know if a model is using MHA, MQA, or GQA. If you look at the config.json file of Llama2, you can tell that it is using basic MHA, because the number of KV heads is equal to the number of attention heads.

"num_attention_heads": 32,
"num_hidden_layers": 32,
"num_key_value_heads": 32,

In the Llama3 config file, in contrast, num_key_value_heads is reduced to 8. This means it is using GQA, where each KV head is shared: 32 divided by 8 = 4 attention heads.

"num_attention_heads": 32,
"num_hidden_layers": 32,
"num_key_value_heads": 8,

Another notable advancement in KV cache efficiency is multihead latent attention (MLA), introduced by DeepSeek, which is the rightmost setup in Figure 6-10. The key difference between MLA and other variants is that, instead of simply reducing the KV counts, It uses a clever way to compress them. DeepSeek’s original paper claims to achieve a KV cache that “is equal to GQA with only 2.25 groups, but its performance is stronger than MHA.”

Kernel fusion and custom attention kernels

Next, let’s look at how to make attention computation even faster by using specialized GPU kernels. Kernels are small, specialized programs that execute on the GPU to perform computations as matrix multiplications, softmax, and other crucial operations for LLMs and many other deep-learning models. Using properly optimized and specialized GPU kernels based on the given model architecture, hardware, and LLM workload can significantly improve your GPU utilization, inference speed, and throughput.

Kernel fusion

One key technique in kernel optimization, widely used in the general machine learning world as well as LLMs, is kernel fusion. Kernel fusion merges multiple individual operations (such as a multiplication operation and an addition operation) into a single one, to minimize the data-movement overhead between memory and compute. In this way, it reuses the data already in the register or shared memory, without needing the round trip of writing back to the GPU global memory and loading again. Figure 6-11 illustrates how kernel fusion works by converting the shapes in one go, without writing and loading back and forth to and from memory in the middle step.

Figure 6-11. Comparing memory and compute interaction with kernel fusion (left) and with kernel fusion (right). Source: He 2022.

FlashAttention

Next, let’s look at one of the attention kernels in detail. FlashAttention represents a major advancement in high-performance attention mechanism computation. The core idea of FlashAttention is that, since attention is bottlenecked by GPU memory bandwidth during frequent reads and writes, making the algorithm hardware-aware (or memory I/O aware) to reduce the I/O combats that memory limitation.

Recall that in Chapter 5, we discussed the hierarchy of GPU memory bandwidth: the main GPU, global high bandwidth memory (HBM), is the biggest within the GPU memory hierarchy, but it also has the slowest data movement. The goal of FlashAttention is to avoid materializing the big matrices in the slow GPU global memory, instead breaking up those large matrices into smaller ones. This is called tiling (or blocking), and it allows all the computation to happen in faster memory, like SRAM or even registers, without needing to perform GPU global memory. Figure 6-12, from the original FlashAttention paper, shows how FlashAttention performs attention QKV matrix multiplication and conversions in tiles iteratively, so all the compute happens in GPU SRAM and only the final output is saved to HBM.

Figure 6-12. How FlashAttention performs QKV computation in SRAM using loops (left) and the performance gain of fused FlashAttention (right). Source: Dao et al. 2022.

The FlashAttention kernel also fuses all attention operations, together with additional key ideas such as online softmax, to achieve one of the best attention kernels available. It is now used by many serving frameworks.

Over time, FlashAttention 2 and 3 have improved on the original. They keep its core ideas but add techniques, such as overlapping General Matrix Multiply (GEMM) calculation with softmax calculation to improve GPU utilization–especially in newer-generation GPUs such as H100.

Kernel optimization is a big topic and an important research area that requires expertise in multiple areas such as GPU architecture, CUDA, performance profiling, and compilers, and is thus well outside the scope of any single chapter. What we want you to take away is this: From a practical and engineering perspective, when serving LLMs, it is important to leverage efficient kernels. Other optimized kernels include FlashInfer, xFormers, and Triton (not to be confused with Triton Inference Server, a totally different concept). They all enable high performance attention for fast model training and serving, exemplifying the power of GPU kernel optimization.

For example, with a few lines of code and environment variables, you can install a custom FlashInfer kernel and its dependencies to run in vLLM:

pip install vllm==0.8.5.post1
pip install flashinfer-python==0.2.2

export VLLM_ATTENTION_BACKEND=FLASHINFER
export VLLM_USE_FLASHINFER_SAMPLER=1
export VLLM_FLASHINFER_FORCE_TENSOR_CORES=1

It’s similar in SGLang, with a tweak of a flag:

--attention-backend {flashinfer|fa3|triton|torch_native|FlashMLA}

Because of the complex and varied natures of different kernels, hardware, and LLM input/outputs, it is relatively hard to give clear-cut advice on which kernel to use. In our experience, finding the appropriate one usually requires some experiments. And many serving backends, such as vLLM and SGLang, have built-in logic to choose which attention kernels to use by default. For example, at the time of writing, SGLang defaults to FlashInfer for non-Hopper machines (like A100 and A40) and to FlashAttention3 for Nvidia Hopper architecture GPUs (like H100, H200, and H20). It is okay in practice to start with the recommended one and pursue other optimization opportunities at first before experimenting with different kernels for additional performance gain.

PagedAttention

One last important way to improve attention is by efficiently managing the KV cache memory. During model serving, the system is constantly creating and storing new KV cache and evicting old cache. The KV cache size grows larger and larger as modern LLMs become capable of longer context lengths. Also, scheduling this large KV cache is never easy, because for LLM serving, the input length can vary by request and the output length is not predetermined. All these unique challenges can cause a lot of inefficiency in GPU memory, because traditional methods would preallocate memory space that usually isn’t fully used, leading to significant memory fragmentation and low usage of real-world memory space.

An important advancement introduced to combat these problems is PagedAttention, along with the popular open source vLLM serving framework that builds on top of it. PagedAttention was inspired by the paging technique in operating systems that divides memory into fixed-size blocks called pages to make it easier to allocate scattered free spaces, and to enable long sequences without requiring contiguous memory. PagedAttention works in a similar way. It first partitions the KV cache into fixed-sized blocks. Then, using a lookup table, it maps query keys to specific blocks. This means the KV cache does not need to be stored in contiguous memory–the blocks can just be accessed individually when needed.

Figure 6-13. One step in the generation process for a request with PagedAttention. A full animation is available from the vLLM blog.

As you can see from Figure 6-13, the full prompt and completion combination is not stored in continuous space in the actual physical memory, pictured on the right. They are stored in three different blocks, starting at block 7, then block 1, and now block 3. Each block contains a maximum of 4 tokens or words, with the last one containing only 2, as it is still actively generating. The block table is the lookup table used to find the physical block where the requested data is located.

The original PagedAttention paper states that, without PagedAttention, “only 20.4% - 38.2% of the KV cache memory is used to store the actual token states,” but that with PagedAttention, it achieves “near-zero waste in KV cache memory.”

Because of these great improvements in reducing memory fragmentation, PagedAttention and its variants have become a basic feature of LLM serving that is enabled for almost all usage, much like continuous batching. You can read much more about the implementation details of PagedAttention in vLLM’s documentation.

Model Compression

LLMs have unlocked astonishing capabilities, but their sheer size causes so many problems in real-world production usage. Acquiring high-performance GPUs can be both costly and challenging. To bring these large and expensive models to consumers, we need to shrink them—smartly.

That’s where model compression techniques come in. These are not crude hacks, but the production strategies that reduce model size and computational load while preserving performance.

Model compression techniques normally fall into three categories:

Quantization

Quantization reduces the precision of a model’s parameter from higher-bit to lower-bit formats, to squeeze more parameters into memory and speed up matrix operations.

Distillation

Distillation transfers knowledge from a large “teacher” model into a smaller and faster “student” that mimics its behavior.

Pruning

Pruning surgically removes redundant weights or attention heads, revealing how much of the model’s capacity is underused.

In this section, we will unpack each of these techniques, show their tradeoffs, and explore practical ways to apply them.

Among the three major model compression techniques, quantization stands out in terms of practicality. It is fast, effective, and typically requires little to no modification of the model training pipeline. These advantages make quantization the go-to choice for compressing and accelerating LLMs in production environments. In fact, quantized models are more widely used in production than you might expect—particularly in scenarios that demand low latency and high throughput, or where models need to run on resource-constrained edge devices.

Given quantization’s impact and prevalence, we’ll spend more time diving into it in this section than the other two techniques.

Quantization

Quantization is the process of reducing the precision of a model’s parameters, such as weights, activations, and KV cache, from high-precision floating-point number formats like 32-bit (FP32) and 16-bit ( FP16 or BF16) to low-bit representation, such as 8-bit (FP8 or INT8) and 4-bit (FP4 or INT4). With quantization, you’re essentially trading lower accuracy in the model data for better serving performance.

Quantization Error

During quantization, you are reducing the precision of the model parameter’s numerical value from high to low. This introduces two types of errors–rounding errors and clamping errors–that can degrade the model’s inference accuracy.

Rounding errors happen when a number (like FP32) can’t be represented exactly in a given numerical format (like INT8), so the system rounds it to the nearest representable value. The difference between the original value and the converted value is the rounding error. For example, let’s say we have a FP32 value 7.6. When we convert it to INT8, since INT8 can’t store the decimal, 7.6 gets rounded to the nearest integer: 8. The rounding error here is thus 8 - 7.6, or 0.4.

Clamping errors happen when a number is too large or too small to fit in the numerical format’s range. In these cases, you can “clamp” the number to a maximum (or minimum) representable value. This happens because formats like FP8 and FP16 have limited exponent ranges. For example, suppose FP8 can only represent values between -448 and +448. If your number is 1,000, it’s too big to present in FP8. So, during FP32 to FP8 compression, you would need to clamp 1,000 to 448 (the maximum allowed value in FP8).

Because clamping can introduce severe distortion—for example, turning 4,096 into 448 in FP8—modern quantization techniques avoid hard clamping when possible. Instead, it’s better to use scaling strategies to better preserve the relative magnitude of values within the target range, even if some rounding error is introduced.

The idea is to apply a scaling factor during quantization to compress the value range of the original data, so that more values fall within the representable range of the lower-precision format (like INT8 or FP8) rather than being abruptly clipped at the extremes. The common scaling strategies we have used are symmetric scaling and asymmetric scaling, as covered on Maarten Grootendorst’s blog A Visual Guide to Quantization.

Now that you have a general understanding of quantization, let’s take a closer look at the different data formats used to store LLM parameters.

Storing numbers

Floating-point (FP) format is a standardized way for computers to represent real numbers, such as 3.14, -42.7, or 1e9. The general structure of a floating-point number is:

Total bits = 1 sign bit + mantissa bits + exponent bits

This represents a number in three parts:

The sign bit (1 bit) indicates whether the number is positive or negative.

The mantissa (or significand) controls the precision (level of detail) of the number’s representation.

The exponent controls scale, determining how large or small the number can get.

There are different floating-point formats, but the most common are FP32, FP16, and BF16. Table 6-1 provides a detailed breakdown.

Table 6-1. Comparing common floating-point number formats and their bit allocations

Precision format
Total bits
Sign bit
Exponent
Mantissa
Approx Range Level

FP32

32

1

8

23

±10³⁸

FP16

16

1

5

10

±10⁻⁵

BF16

16

1

8

7

±10³⁸

FP32 is considered full precision and uses a total of 32 bits: 1 sign bit, 8 bits for exponent and 23 for mantissa. Its wide value ranges and good precision, as shown in Table 6-1, make FP32 the standard format for training deep-learning models.

The FP16 and BF16 formats are both half the size of FP32, at 16 bits. The tradeoff made between FP16 and BF16 is that using more bits for exponent or for mantissa results in getting better precision or a better value range, respectively.

Compared to FP16, BF16 allocates more bits to the exponent (8 bits, versus 5 in FP16), which trades off precision (fewer mantissa bits) for a significantly wider numerical range. Interestingly, both BF16 and FP32 use 8 bits for the exponent, meaning they share the same value range. This makes conversion from FP32 to BF16 more straightforward than from FP32 to FP16—there’s no risk of clamping due to range limits, only a loss in precision due to the reduced mantissa. This design makes BF16 particularly well-suited for deep learning training workloads, where dynamic range is more critical than fine-grained precision, while also reducing memory and compute overhead. Figure 6-14 is a visual comparison of FP32, BF16, and FP16 formats.

Figure 6-14. Visualization of FP32, BF16 and FP16 precision and range. Source: Maarten Grootendorst

In the LLM serving scheme nowadays, FP32 precision is less used, and most model checkpoints are provided either in FP16 or BF16.

Now that you’ve seen the full-size data format, let’s examine the quantized number formats. There are two types of quantized number representation: integer-based representation, such as INT8 and INT4, and floating-point-based representation, such as FP8 and FP4, which have become more popular even while we’ve been writing this book.

The key difference between these two formats is how they distribute data points. In the integer-based format, the data distribution is completely uniform, meaning the data points are evenly spaced across the whole value range and the number precision does not change. In contrast, floating-point has nonuniform distribution: that is, there are more data points around zero and less data points at the extremely large or small numbers.

The difference in data density between the FP and integer-based formats is because FP number values are logarithmic. Float number’s value (in FP8 format) is:

Value = (–1)sign × (1 + mantissa) × 2exponent – bias

As you increase the exponent, the distance between representable numbers also increases. For example, between 1.0 and 2.0, you can represent many values (like 1.01, 1.001, and so on), but between 1,000,000 and 2,000,000, the smallest step size becomes much larger, like 128 or more (see Figure 6-15).

Figure 6-15. Data point distribution in INT8 and FP8 formats.

Depending on the data distribution of a given model’s parameters (for example, model weights and activation values), you may want to choose one format over the other to retain the best accuracy.

How does quantization help in model serving?

It is very important to understand why and how quantization helps increase model serving throughput and reduce latency, so that you can choose the best quantization method for your use case.

The first reason is data size. Whether you need 8 or 16 bits directly impacts the model size, on disk and in GPU memory. Recall (from Table 5-6 in Chapter 5) that FP16 precision corresponds to 2 bytes (1 bytes = 8 bits) per parameter. Here’s how you would calculate the data size of a 7-billion-parameter model with FP16 precision:

~7 billion parameters * 2 bytes/parameter = 14 billion bytes = 14 GB

However, quantizing this model to INT8 immediately cuts its size in half, to 7 GB. This is a huge win when you don’t have a big enough GPU memory to fit the unquantized model.

Similarly, when serving a large model, quantizing it can mean fitting it in one machine or node instead of needing multiple nodes and dealing with cross-node communications. The reduced memory footprint also frees up space for the KV cache, allowing the system to handle more concurrent requests and boosting its overall throughput.

You also learned in Chapter 5 that GPU memory bandwidth is the key to LLM performance optimization, especially in the Decode phase, since data movement is the bottleneck. Quantizing the model and making it smaller directly reduces the data movement needed, which greatly helps to reduce model inference latency.

Lastly, lower precision can also mean faster computation. You looked at a similar table In Chapter 5, Table 5-2, when we compared GPU hardware. Now look closely at Table 6-2, which compares FLOPS across different precisions. It is apparent that when we reduce the bits by half, from 16 bits to 8 bits, we typically can double our FLOPS; the same goes for quantizing weights to 4 bits.

Table 6-2. H100 FLOPs across different precisions

H100 SXM

FP64

34 teraFLOPS

FP64 Tensor Core

67 teraFLOPS

FP32

67 teraFLOPS

TF32 Tensor Core

989 teraFLOPS

BFLOAT16 Tensor Core

1979 teraFLOPS

FP16 Tensor Core

1979 teraFLOPS

FP8 Tensor Core

3958 teraFLOPS

INT8 Tensor Core

3958 TOPS

In short, quantization can help reduce model size, reduce data movement during computation, and speed up computation.

Weight-only versus weight-and-activation quantization strategies

Quantizing on the model weight only and quantizing on both the weight and activation are two of the most common quantization strategies. You might have seen some quantization notations like W4A16 and W8A8: this type of notation is used to describe the bit-widths used for a model’s weights (represented by W) and activation parameters (A), which correspond to the model’s middle-step inputs and outputs. For example, W4A16 means 4-bit model weights and 16-bit activations, and W8A8 means 8-bit model weights and 8-bit activations

Weight-only quantization means that the model weights are quantized but the activation parameters are not. This reduces the model size but, during execution, the lower-bit values are dequantized back to high-bit formats. Thus you benefit from the reduction in model size and data movement, but do not achieve faster computation. In fact, because you need to dequantize the weight before computation, it adds a little overhead.

To avoid dequantization in model inference, you can use mix-precision kernels, such as the Marlin kernel for Ampere-generation Nvidia GPUs such as A100, and the Machete kernel for Hopper-generation Nvidia GPUs, such as H100. These kernels can efficiently handle matrix multiplications by fusing the quantized weight access. They can perform multiplication between, say, an INT4 matrix and a FP16 matrix in one pass without dequantization. In practice, many people turn these kernels on by default when serving weight-only quantized LLMs in LLM serving engines.

Weight and activation quantization not only reduces model size and data movement; because it quantifies activation, it also helps achieve higher FLOPs with a compute-bound workload. However, activating it is more complex than with weights-only because complexity depends on both model weight and model input, which changes all the time. This means that, when quantizing activation, you need to choose when to do the scaling. You have two options. With dynamic scaling, you’re basically calculating the scaling factor on the fly as you do the inference. In static scaling, the scaling factor is estimated before the model is deployed and precalculated using a calibration dataset. The tradeoff is that, while dynamic scaling gets better accuracy, it is less performant than static scaling because it calculates the scaling factor during the actual inference.

The most common weight-only quantization used at scale in production is W4A16, using the GPTQ or AWQ quantization methods. The most common weight-and-activation quantization widely used in production at the time of writing is W8A8 (which was previously in INT8 format and recently moved to FP8 format). Table 6-3 compares the benefits of adopting quantization on W4A16 versus on W8A8.

Table 6-3. Comparing W4A16 and W8A8 quantization

W4A16
W8A8

Model size reduction/data movement

75% (original size/4)

50% (original size/2)

Compute FLOPs

No change

2x

Prefill phase (compute bound)

No change

Improved

Decode phase (memory bandwidth bound)

Improved (better at low batch)

Improved (better at high batch)

Best for

Long generation, latency sensitive, low batch

Long context, high throughput, high batch

W4A16 reduces the model footprint twice as much as W8A8, making it occupy less GPU memory and require less GPU memory bandwidth, so it greatly helps with the Decode phase. Comparatively, W8A8 excels in compute-bound workloads such as long-context Prefill phases and high-batch, high-throughput scenarios where the Decode phase is batched and pushed towards being compute bound.

How to choose a quantization strategy

In real-world production deployment, if the model is large and requires 4x model size compression to fit inside a single GPU to avoid cross-GPU communication, or GPUs within a single machine to avoid cross-node communication, W4A16 is the way to go. On the other hand, the benefit of quantizing activation in more efficient compute usually outweighs the memory saving in high batch, moving the bottleneck of the model execution from memory bandwidth bound to compute bound. For production workloads, if we can achieve a certain latency SLA with W8A8, not needing W4A16, we try to push for higher effective batch and high throughput per model instance to reduce cost.

Now, let’s revisit the data format. In the past, W8A8 quantization was mostly done using the INT8 format, where both weights and activations were clamped to the fixed ±127 integer range for simplicity and hardware compatibility. Recently, a lot more models have been deployed to production in FP8 variants (such as E4M3 and E5M2) instead. In a 2022 paper, Nvidia introduced the FP8E4M3 format and showed how it can improve serving performance while maintaining minimal accuracy loss compared to FP16, without needing calibration to maintain model accuracy (like INT8 does).

Among the two variants of FP8, E4M3 is more commonly used for inference because it retains more precision than the E5M2 variant, despite E4M3 having smaller dynamic range and thus could need scaling. Table 6-4 compares the two.

Table 6-4. Floating-point precision formats and their bit allocations

Precision format
Total bits
Sign bit
Exponent
Mantissa

FP8 (E4M3)

8

1

4

3

FP8 (E5M2)

8

1

5

2

FP8 is not always applicable

Another thing to consider is hardware compatibility: not all GPUs or accelerators support FP8 format. Of the Nvidia GPUs, only Hopper and Blackwell support FP8. This means that if you have only A100s or older GPUs, running FP8 might not achieve the expected full performance gain.

Because this is a more engineering-focused book, we will not go deep into the quantization process. For more information, we recommend Grotendorst’s A Visual Guide to Quantization.

Hands-On Quantization

In this section, let’s try some examples to see how to apply quantization in practice.

In the Google Colab notebook, you can find the code that hosts Qwen/Qwen2.5-7B-Instruct model in three different variations: the original, GPTQ W4A16 quantized, and W8A8 quantized in a vLLM server. For each model, we’ll show you how to conduct benchmark tests with different concurrencies. You’ll record their TTFT, TPOT/ITL, and throughput to better understand the performance gain and tradeoffs between them.

Find the quantized model

In the notebook, we execute both the original model and quantized versions of the models as a standalone server. To execute a model, you’ll first need to either save it to a local or cloud location, or directly reference a model on Hugging Face. In many cases, the popular foundation models and their quantized versions can be found directly in Hugging Face so you might not have to quantize the model yourself. For example, Hugging Face provides all the quantized variants of the model uploaded already, on the right-hand side, as shown in Figure 6-16.

Figure 6-16. Quantized model variants available on Hugging Face.

Quantize a model yourself

In this case, you can simply pass the model ID, “Qwen/Qwen2.5-7B-Instruct,” as an argument to the server startup command. You can also replace the model ID with a model path if needed.

hf_model_id = "Qwen/Qwen2.5-7B-Instruct"
!vllm serve {hf_model_id}

If you want to use a different calibration set to quantize a foundation model, or quantize a custom fine-tuned model yourself, there are many help libraries, such as GPTQModel, AutoAWQ, and LLMCompressor.

Here’s a code snippet for GPTQ quantization:

from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
dataset = ["Gptq is an easy-to-use model quantization library with user-friendly apis, based on GPTQ algorithm."] # calibration dataset
gptq_config = GPTQConfig(bits=4, dataset=dataset, tokenizer=tokenizer)

quantized_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B-Instruct", device_map="auto", quantization_config=gptq_config)

Running benchmarks

Once the vLLM server is up, you can start sending it requests and observe how models quantized using different methods perform. To run the benchmark, use the benchmark script provided inside vLLM’s repo.

Performance analysis

After running the benchmark scripts for all three model variants, we gathered the performance figures and plotted them, as shown in Figures 6-17, 6-18, and 6-19.

Figure 6-17. Comparing TTFT between the original model and models quantized with GPTQ W4A16 and FP8 W8A8 on Qwen2.5 under different request concurrency.

You can see that in low-concurrency settings, GPTQ W4A16 provides the best performance, with about a 300% performance gain in latency and throughput over the original model. FP8 W8A8 also provides a decent gain, but 150%. In low-batch settings, the bottleneck is on GPU memory bandwidth. GPTQ reduces the weight to INT4, which is four times smaller, and that really shines here. For use cases where you need the best latency, such as chatbots or AI agent flows (where multiple calls to the model are needed and the aggregated latency reduction is crucial), W4A16 weight-only quantization is the way to go.

Figure 6-18. Comparing TPOT/ITL between the original model, GPTQ W4A16, and FP8 W8A8 on Qwen2.5 under different request concurrency.

Next, let’s see what happens as concurrency goes up along the x-axis from the left to the right side of each chart, in high-batch processing. GPTQ W4A16 shows its weakness here because, as a weight-only quantization method, its computation is still in 16-bit. The TTFT increase (Figure 6-17) is most apparent with GPTQ W4A16, which is even slower than the original model. This is because it does not save anything in computation, and incurs dequantization overhead. FP8 W8A8, on the other hand, performs much better now because the activations are also quantized to FP8, resulting in faster computation. If W8A8 latency is already good enough to satisfy SLA requirements for your use case, this is the time to push more batches, so that each model instance can serve more requests in parallel to save cost.

Figure 6-19. Comparing request throughput between the original model, GPTQ W4A16, and FP8 W8A8 on Qwen2.5 under different request concurrency.

Other quantization methods

So far, when we’ve talked about quantizing the weights and activations, we are mostly focusing on the FeedForward layers, not the attention layers. Even though FeedForward layers comprise the majority of the parameters and computation, the attention mechanism and KV Cache quantization are also applied in certain cases.

KV Cache quantization and attention quantization

The first step is quantizing the KV cache. In both weight-only and weight-and-activation quantization, the KV cache is usually left in a high-precision format to reduce its size and free up more GPU memory. With more GPU memory, you can increase the batch size to gain higher throughput. Quantizing the KV cache can also benefit techniques such as prefix caching, allowing more prefixes to be cached without recalculating or reloaded. (We will discuss prefix caching in detail later in this chapter.)

However, quantizing KV cache usually will not significantly reduce latency. This is because, if the attention calculation is still using a high-precision format, quantizing the KV cache only makes it smaller. That doesn’t mean you can gain more FLOPs during the calculation. During the attention calculation, similar to weight-only quantization, quantized KV cache also needs to be dequantized back. Thus, to fully reap the benefit of a quantized KV cache, you need to use it in conjunction with quantized attention kernels.

Overall, we usually still start with weight-only or weight-and-activation quantization. For example, we start with FP8 with both weight and activation. Then to tackle long context in limited GPU memory or to gain more throughput for decoding a heavy workload, we add FP8 KV cache quantization and enable an FP8 attention kernel when needed.

GGUF Quantization

GPT-Generated Unified Format (GGUF) is a very different type of quantization. GGUF is popular among users seeking local, portable, and low-resource deployments because it focuses on running LLMs on CPU and) or Apple Silicon (Metal) instead of GPU, but offloads some parts to GPU when needed, if available. It also provides a lot of different quantization for levels and formats to meet your needs when high-end GPUs are not available. In Chapter 8, we will discuss deploying GGUF quantized models.

Accuracy Tradeoffs and Mitigation

So far we have examined many different quantization techniques and their tradeoffs. However, the biggest tradeoff when applying quantization is between model accuracy and serving performance (throughput and latency). If a quantized model can’t maintain the accuracy of the original model to sustain enough quality, all the performance gain cannot be materialized because the inferior product can never be deployed to production. Luckily, there’s been a lot of testing and research on maintaining model accuracy. Nowadays, methods such as GPTQ W4A16, AWQ, and FP8 quantization on both weights and activation have all shown minimal loss in real-world accuracy metrics, as shown in Figure 6-20. On the other hand, the serving performance gain from quantization sometimes enables you to run a quantized version of a bigger model. For example, you could deploy a 12-billion-parameter model in FP8 over an 8-billion-parameter model in FP16 to achieve comparable latency, higher throughput, and better model accuracy.

Figure 6-20. OpenLLM Leaderboard scores for different quantization methods with different model sizes, by Neural Magic. Source: Red Hat.

This research has also shown some patterns in how different quantization methods affect models of different sizes and context lengths, which can guide you in finding the best quantization level and method to retain satisfactory accuracy. For example, you could try to push the boundary with more aggressive weight-only quantization and long context serving; the larger the model is, the more sensitive its accuracy will be to model weight and KV cache quantization. KV cache quantization is relatively less intrusive in terms of accuracy, but doesn’t yield a significant performance again. If you’re struggling with KV cache space and ITL, it is still a good method, but less of a priority.

Current research is moving towards even lower bits such as FP4, FP6, and W4A8 quantization with newer Blackwell generation Nvidia GPUs. So far, it has been quite hard to maintain good model accuracy, but thanks to per‑tensor and per‑channel scaling, outlier‐aware clipping, and newer techniques all working together to minimize quantization error, expect to see these lower-bit options used in production in the near future.

One last way of mitigating accuracy issues is performing quantization-aware training. All of the quantization techniques we’ve described so far focus on post-training quantization (PTA), which doesn’t require access to the training pipeline at all. This enables you to quantize the foundation and custom-trained models you have access to but did not necessarily train yourself. Quantization-aware training (QAT), on the contrary, performs fake quantization during the training phase to better preserve accuracy. The idea is that the fake quantization lets the model adapt to the quantized weights during training, providing better accuracy than PTA. In practice, though, we’ve found that PTA is a lot more popular than quantization-aware training because it’s easier to use and still decently accurate.

After quantization, you can evaluate the model’s accuracy to make sure it is still satisfactory. A commonly used tool is LM Eval, which supports a wide variety of setups.

Here is a code snippet for accuracy evaluation:

lm_eval --model {hf|vllm|sglang} \
    --model_args Qwen/Qwen2-7B-Instruct \
    --tasks gsm8k_cot \
    --device cuda:0 \
    --batch_size auto

Distillation

Now let’s look into another technique in model compression, one that we believe has the biggest potential of the three to improve latency and throughput: model distillation. Model distillation is unique compared to quantization and pruning (discussed in the next section) in that, instead of trying to reduce the size of the original model, it basically trains a new small model. It tries to transfer the knowledge encoded in the large original “teacher” model into a smaller “student” model. In that process, the student model mimics the teacher model.

Figure 6-21. Model distillation pipeline to train a student model from a teacher model.

Figure 6-21 shows how a teacher model can generate outputs that can be used to train a much smaller student model. The outputs from the teacher model aren’t necessarily limited to just the output token (which is usually called hard label), but also can include logits for the probability distribution of the prediction and losses. Thus, to create a distilled student model, you need full access to the original teacher model, not just API calls to get the output tokens.

Let’s take DeepSeek’s distilled models as an example. The original DeepSeek R1 model has 671 billion parameters and a Mixture of Experts (MOE) architecture (we will discuss MOE models in Chapter 7). DeepSeek has also released multiple models distilled from the LLama and Qwen families of open-source dense models, ranging from 1.5 billion to 70 billion parameters. From a serving perspective, this reduces model size by 10 times or more and can mean huge improvements in serving latency and throughput. Table 6-5 compares the original with a 70-billion-parameter distilled model on the benchmarks.

Table 6-5. Comparing evaluation benchmarks between the original DeepSeek R1 and its distilled 70-billion-parameter model. Source: DeepSeek.

Benchmark
DeepSeek-R1-671B
DeepSeek-R1-Distill-Llama-70B

MATH-500 pass@1

97.3

94.5

GPQA Diamond pass@1

71.5

65.2

LiveCodeBench (Pass@1)

65.9

57.5

Table 6-6 shows the pros and cons of quantization and distillation. If you’re trying to figure out whether to use a distilled model or distill your own model to improve serving performance, our general guidance is to start with an easy, low-cost solution first. Let’s say you are considering deploying either DeepSeek-R1-671B orDeepSeek-R1-Distill-Llama-70B, If the distilled model is already available, you can start by evaluating it to see if it meets your accuracy benchmark. If it is good enough, go with the distilled model. You can perform quantization on the distilled model to further improve serving performance and save cost. However, if you do not have a distilled model ready to use, which is likely to happen most of the time in real life, quantization should be your first step, given the high cost of performing distillation and the bigger accuracy loss.

Table 6-6. Comparing quantization and distillation.

Quantization
Distillation

Accuracy drop

Low (usually 3% or less)

Much higher than quantization

Speed gain

1.5x to 3x

Much more, with a tradeoff in accuracy

Ease of use

Very easy for post-training quantization; you only need the original model weight

Harder if the distilled model is not readily available; training costs can be as much as 10% of the original model cost. Usually done by researchers who train the original model.

Pruning

The last model compression idea we’ll look at, model pruning, is also the least popular, since (as we write this in mid-2025) more work and research are still needed to make it work for production. The core idea of model pruning is that models are usually overparameterized; pruning their redundancies can achieve better compression and thus improve serving performance.

Pruning is usually divided into structured and unstructured: Structured pruning removes sections from the model entirely, while unstructured pruning removes individual weights with more flexibility. Notably, in the middle lies semi-structured sparsity, also called 2:4 structured sparsity, where 2 out of every 4 elements are pruned. One semi-structured sparsity model is Sparse LLama 3.1 by Neural Magic, which claims to have achieved 98% accuracy recovery, 30% higher throughput, and 20% lower latency from sparsity alone with vLLM.

Figure 6-22. A2:4 structured sparsity pattern and compression. Source: Nvidia.

In Figure 6-21, for every group of four consecutive values in the original matrix on the left, two values are zeroed out (shown in white). This forms a sparsity rate of 2:4, or 50%. Because Nvidia’s newer GPU architectures, such as Ampere and Hopper, have sparse tensor cores that accelerate this type of structured sparsity, 50% sparsity can yield double the speed in matrix multiplication.

Next, we’ll finish up the chapter with a look at prefix caching.

Prefix Caching

In software engineering, caching is a very common technique for storing frequently accessed data in fast, possibly temporary storage to reduce latency and improve overall performance. In machine learning, caching the model prediction outputs of the same requests is very common in production, done on the client side, the server side, or both. The client side can use caching to avoid sending unnecessary duplicate requests, while the server side can use it to return the cached output without executing the model, to reduce computation.

An example of response caching on the server side is in Triton Inference Server, where each inference request (including model name, model version, and model inputs like tensor name, shape, datatype, and data) are hashed and cached along with the model output as key-value pairs. When a new request comes in, if the hash is found, the inference output is extracted directly from the cache.

The main reason request caching works is that storage is a lot cheaper than compute, and thus accessing storage is a lot faster than recomputing the result. However, while a general request-caching solution will help ML inference, the gain in LLM serving is usually small because LLMs take free-form human-written text as input. People can ask the same question in so many different ways that the cache hit rate of the naive request hashing can be very low.

A widely used, proven, powerful caching technique in LLM serving is prefix caching. Instead of matching the entire prompt, prefix caching tries to match the “prefix” of a prompt to those of all of the other prompts the model has processed before. If there is a match, the KV cache of the prefix part that is matched no longer requires recomputation. Instead, the cached data is retrieved–in simple cases, directly from GPU memory.

The KV cache can also be offloaded to other places, such as the CPU, local SSD, or distributed across GPUs; on completely different machines; or even on distributed external storage systems–but we will discuss more complex storage design in Chapter 7. For now, though, let’s assume the KV cache is all saved in current GPU memory. In this simpler case, if prefix caching is not enabled, after each request is completed, all KV cache generated for that request is removed from GPU memory. If prefix caching is enabled, the requests’ KV caches are kept in GPU memory as long as there is still space available. In most prefix-caching implementations, a Least Recently Used (LRU) mechanism is used to evict the KV cache when not enough GPU memory is available.

RadixAttention

One of the prominent prefix-caching solutions is RadixAttention, which was introduced along with the SGLang serving framework. RadixAttention leverages a radix tree, which is a data structure that works like a trie (also known as a prefix tree): a form of string-indexed lookup data structure to keep track of the prefix strings for the KV caches. When the tree grows too large, it also uses an LRU mechanism to evict the KV cache, applied recursively on the tree’s leaf nodes.

Figure 6-23. A radix tree example for two requests sharing the same prefix. Source: Zheng et al. 2024.

Figure 6-22 shows a very basic example of two requests that share the same head node, You are a helpful assistant, and then branches afterwards. When two requests share the same prefix, RadixAttention can find corresponding nodes in the tree structure, which is stored in CPU memory. With each node mapping to the KV cache, which is stored in GPU memory, it can reuse anything in the KV cache without incurring recomputation. As new requests come in over time, the tree grows by adding leaf nodes, and gets trimmed by LRU when it becomes too large.

Use Cases

Now let’s examine in what cases prefix cache works and in what scenario it becomes useful. Let’s start with an example. Consider these three prompts:

Prompt 1: Hi, what is the weather like today?
Prompt 2: Hi, what is the weather like now?
Prompt 3: What is the weather like today?

Prompt 2 can reuse prompt 1’s prefix KV cache from “Hi, what is the weather like…?” However, prompt 3, because it does not have the “Hi” as prefix, would need a complete fresh recalculation. You might wonder–does that mean it’s not really very powerful? But in real-world LLM production, it is already proving very helpful indeed. Here are two major scenarios where prefix caching works at its best.

Scenario 1: Multiturn chat

One user interacts with an LLM, going back and forth multiple times. The prior chat history is not omitted when the user sends a new prompt. Instead, it is kept and concatenated with the new prompt the user sends. For example, the user first asks:

User: What is the weather like today?
LLM: The weather is 70F, but with a likelihood of rain.

The user sends a second question:

User: Could you tell me the actual possibility?

The actual prompt sent to the LLM is not just that question alone. The LLM is usually given all the prior interactions as context, so that it can produce the best response possible.

The LLM knows that the user is not asking about an arbitrary possibility, but about the possibility of rain today.

Without prefix caching, the LLM would perform the Prefill phrase for the first two lines again when the user asks the second question. With prefix caching, since the LLM has already processed the first two lines, they are kept in the KV cache in GPU memory instead of getting discarded. The LLM sees the match, knows that it has processed the first two lines before, and reuses the stored KV cache, without starting over from the very beginning.

This might look like a small reduction at first, but as the context of the multiturn chat grows longer and longer, without prefix caching, the user would experience longer and longer wait times when they ask additional questions. Most notably, the TTFT, which dominates the customer experience in LLM chatbots, would increase.

Scenario 2: Long context serving

Another scenario where prefix caching provides critical performance gain is long context serving. LLMs’ context size continues to grow, from 4,000 to 128,000 and even 1 million. This gives users a lot of opportunity to feed as much as relevant information into the input prompt, without always relying on extracting relevant information with RAG. But a very long context inside the prompt poses challenges in the Prefill phase of LLM serving, which can make the TTFT intolerably long. Here comes prefix caching to the rescue: it saves the KV cache of all of the relevant information to avoid recomputation, as long as the user carefully reconstructs the prompt to ensure that its prefixes do not vary. This can result in cache hits every time.

Prefix caching is now enabled by default for almost all scenarios, even when the cache hit rate and TTFT improvement rate are not as good as the two mentioned above. This is because modern optimized serving engines are capable of achieving minimal or even zero overhead when turning prefix caching on. In other words, there is almost no drawback to using it. Even if the cache hit rate is only 5% in your use case, the blazing TTFT speed of that 5% of the customer experience is still worthwhile.

Best practices

To best utilize prefix caching, your key target should be to improve the cache hit rate. The first and foremost way to do that is to carefully structure prompts so that the LLM cleanly separates the static parts (which target the prefix hit at the very front of the prompt) from the dynamic parts (which correspond to user input at the back as the suffix).

You can use the following template to construct prompts so that everything in the system and context results in cache hits:

<system>
You are a helpful assistant.
<context>
Document: {puts the relevant static context here each time}
<user>
{dynamic questions from user}

Even just changing the word Document to Documents –just one extra letter–can result in a cache miss. Thus, assembling the prompt programmatically and consistently is very important.

The same is true for RAG use cases. Even though these cannot have as high of a cache hit rate as the long-context scenario, you still want to achieve a high cache hit rate and make the hit as long as possible. For example, in RAG, your context would be shaped like this:

Document 1: <retrieved_text_chunk_1>
Document 2: <retrieved_text_chunk_2>
Document 3: <retrieved_text_chunk_3>
Document 4: <retrieved_text_chunk_4>

Or like this:

Document 1: <retrieved_text_chunk_1>
Document 2: <retrieved_text_chunk_2>
Document 3: <retrieved_text_chunk_5>
Document 4: <retrieved_text_chunk_7>

It is still important to maintain the same formatting and structure each time. Moreover, consistent ranking and proper deduplication can also help to achieve the longest prefix hits. For example, even if you can’t get a prefix hit of the full context, you can still achieve a prefix cache hit all the way to “…Document 3:” between the two context examples above, which is still a win.

Scaling Prefix Cache

As your serving traffic increases with more requests and more customers, it is inevitable that you’ll need to scale your setup.

First, let’s consider a setup with only a single model instance. In Chapter 5, you saw in Figure 5-11 that KV cache can take up a significant percentage of GPU memory during LLM serving. With prefix caching, that’s even more true, since it is advantageous to cache as many requests as possible to increase the prefix cache hit rate. Thus, reserving enough space in the GPU to cache common prefixes is important. (See Chapter 7 for more on other storage methods.)

Another unique challenge when using the prefix KV cache is that, as traffic increases, your setup will need to scale horizontally, to have more model instances working in parallel. To route and distribute requests efficiently to all the model instances, common load-balancing techniques are round robin, where requests get distributed across the servers in a sequential or rotational manner; least connection, where requests get distributed across the servers with the fewest active connections; or based on custom metrics, such as GPU utilization for ML workloads. With prefix caching, since the prefix KV cache is local, it becomes important to have a smart routing layer. Like consistent hashing, it creates an affinity between prefixes and mode instances, so that requests are routed to the model instance (or instances) that has the prefix KV cache already stored, rather than to a new instance to start the Prefill process over. Figure 6-23 illustrates this setup.

Figure 6-24. A simplified illustration of cache-aware routing, where the router routes requests to achieve high cache hit rates, based on the model instance’s request prefixes and the prefix cache.

With the routing layer, not all model instances need to cache all of the prefixes as well. This helps reduce the amount of GPU memory required and avoids constant evictions and recomputation. But what if the prefixes are very long or you need to cache a lot of different prefixes, and GPU memory still isn’t enough? One solution is to offload the KV cache to CPU memory and/or SSD when needed (see Chapter 7).

Finally, the LLMs you are serving will usually be multitenant: in other words, different customers will share the same model endpoint and instance. It would be too expensive to create one model instance for each customer, and batching different customers’ requests together can greatly reduce the real-world batch size to improve GPU utilization and save cost. But, with prefix caching, customer A’s prefixes might be the same as customer B’s prefixes by accident. Customer A could then perform enumerations to infer the data processed by another customer, based on the latency it observes. This is a behavior we do not want.

One way to isolate prefixes to a single tenant is to inject a unique ID for each customer into the prompt, as Nvidia discusses in its technical blog. In the example below, a user ID or session ID is added between the system prompt and the context section, so prompts from different customers can only share the same system prompt when they have different IDs.

<system>
You are a helpful assistant.
<id> {user id or session id}
<context>
Document: { puts the relevant static context here each time}
<user>
{dynamic questions from user}

Summary

In this chapter, we delved into a spectrum of techniques to optimize the LLM serving performance in real-world deployment settings.

The chapter began by exploring online request-batching and scheduling techniques. By aggregating multiple incoming requests into batches, models can process them more efficiently, optimizing arithmetic intensity and overall throughput. Dynamic batching is a common strategy for general ML workloads. A more advanced technique, continuous batching, which addresses the LLM specifically, takes precedence over dynamic batching to continuously add new requests into the batch, reducing GPU idle time. Today, continuous batching is a default solution for LLM online serving. On top of that, to better balance TTFT and ITL and improve overall throughput, you can enable and tune chunked prefill to tackle very long contexts and interactive use cases.

Next, we looked at how to improve attention calculation. One way of doing that is to reduce KV cache size by sharing the same KV across multiple heads while maintaining accuracy in the evolution from MHA to MQA to GQA. Another is by performing smart compressions, such as MLA in DeepSeek. Then we explained how kernel optimization can improve attention computation performance, for example by doing kernel fusion to reduce round trips back to memory during calculation. We also looked at FlashAttention where, with techniques such as tiling, the algorithm smartly avoids using HBM. Instead, it keeps data in SRAM to avoid HBM bottlenecks based on hardware IO specs and limitations. Finally, PagedAttention, which is relatively orthogonal to kernels for computation, focuses on optimizing GPU memory access by partitioning the KV cache into small blocks called pages to reduce memory fragmentation. This is now the default choice in KV cache management for LLM serving.

To further improve model serving performance, model compression trades some accuracy for significant size reduction and execution performance improvement. The most common technique used in industry is quantization–more specifically, post-training quantization–where after model training completes, model weights are quantized from high-bit to low-bit formats. We discussed quantization techniques and low-bit formats, along with their tradeoffs. We also glanced at other model compression techniques, such as model distillation, where a large model can teach a small model, and model pruning, to reduce unnecessary and insignificant weights in the model.

Finally, we showed you how prefix caching, as an LLM-specific request-caching method, helps reuse partially seen requests, trading storage for less computation and very fast TTFT. Prefix caching excels at multiturn chat and long-context scenarios, which are becoming more and more essential in modern LLM applications.

Now you should be equipped with the essential knowledge to serve LLMs efficiently, especially smaller LLMs that run on a single GPU. In the next chapter, we will dive into advanced topics to help you run large LLMs in distributed fashion, and improve the LLM serving as a system over focusing on a single model replica.

Chapter 7. Advanced LLM Optimization Techniques

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 7th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

After the last chapter, you are equipped with essential techniques to combat many of the challenges of LLM serving optimization, especially those that are not overly large and fit into one GPU. For larger LLMs with, for example, more than 100 billion parameters, one GPU is usually not enough to load the model to GPU memory and generate at a satisfactory latency. In this chapter, we explore advanced techniques to further enhance LLM serving performance, including:

Speculative decoding to speed up the Decode phase of LLM generation for faster ITL

Multi-GPU and multinode serving for large LLMs that do not fit or are not performant enough when running on a single GPU

Prefill–Decode (PD) disaggregation to decouple the Prefill and Decode phases and thus maximize resource utilization, slash tail latency, and fine-tune the cost-versus-performance tradeoffs of Prefill and Decode independently

Advanced KV caching techniques to achieve lightning fast TTFT and a high cache hit rate

Speculative decoding

What if a single technique could singlehandedly improve latency–especially ITL–by a factor of two or three? Meet speculative decoding, a novel approach that is particularly useful for long, reasoning-heavy generations.

In a large ML system with, say, a million data points for retrieval or recommendation, it’s common to use a small but less accurate model to do the first round of filtering. Once you’ve significantly reduced the number of candidate data points to around 1,000, you then apply a larger, more accurate model on those remaining candidates to produce the final result.

Speculative decoding does something very similar, but at the token level. It leverages a small model, which we call the “draft model,” to help generate candidate tokens for the large “target model.” This speeds up the Decode phase of LLM generation. In other words, during token generation, we have our small draft model speculate on what the next tokens could be. Rather than generating tokens one by one, the target model now verifies the generated tokens. If the target model finds that the draft model is doing a decent job, it accepts the token and jumps ahead.

Detailed Steps

We’ll start with one iteration of speculative decoding, where the draft model performs token speculation and the target model then perform parallel verification, as shown in Figure 7-1.

Figure 7-1. One iteration of speculative decoding dissected into four separate steps

First, the draft model quickly generates K tokens. K here represents the number of tokens the draft model speculates, which is an important parameter to tune for the optimal value. Immediately, the target model performs its forward pass, during which it verifies if the drafted tokens are good or not.

In token generation, the next token is generated with a probability value. For example, given the phrase The soccer team of the United, the draft model might speculate that the next token will be States, with a probability of 0.6, whereas the probability of Kingdom is 0.3 and that of Nations is 0.1. If our target model thinks States is 0.8 probable, which is higher than 0.6, it will accept the drafted token. If it thinks States is 0.4 probable, usually we calculate it probabilistically by dividing 0.4 by 0.6.

If the target model still accepts the draft token, it can move forward to verify the next token. If it decides to reject the draft token, we discard all the following draft tokens. We do this because the generation is autoregressive, so none of the following tokens will be relevant. The candidate model will instead generate one token that it thinks appropriate. This is represented by Token3 in Figure 7-1.

Importantly, this process does not compromise accuracy. In other words, the final output of the speculative decoding process should be identical to what a nonspeculative autoregressive model would produce. This is achieved by the verification step in particular. When a token is rejected, the target model corrects the problem by sampling a new token from its own modified distribution, which it adjusts to exclude the rejected token. The process restarts from this new, confirmed point.¹

Tuning and Usage

This section provides some important tips for applying speculative decoding in real-world use cases.

Choose an existing small model

To maximize the performance of speculative decoding, the key is to use a small model that maximizes both speed and acceptance rate. A model that is from the same family means the same tokenizer, and pretraining can really help in achieving a high acceptance rate. It’s also a good idea to be aggressive in quantizing your draft model, since you always have the target model as the fallback.

If you can perform some training, distilling a small draft model from your target model is usually a better choice than picking an existing small model. This alignment reduces the mismatch between the drafter’s predictions and the target’s verification step, which in turn increases the acceptance rate of drafted tokens and improves overall efficiency. In practice, this makes speculative decoding more robust, since it specializes the small model for the target’s style and domain, rather than relying on general pretraining alone.

Use self-drafting

To push the performance further, self-drafting tools like Medusa and Extrapolation Algorithm for Greater Language-model Efficiency (EAGLE) modify the target model with lightweight prediction heads and early-exit paths to produce multistep lookahead inside the same model. This means we no longer need the extra draft model, saving GPU memory, simplifying deployment, and achieving good alignment with the target model.

Figure 7-2. Medusa uses multiple heads to predict multiple subsequent tokens in parallel (Source: Cai et al., 2024).

Figure 7-2 shows an example of self-drafting with Medusa. On the left, a forward pass with the original model only generates candidates for the first token, shown as It, I, As on the top right. What’s new is that the Medusa heads in the middle perform extra generations of the second, third, and fourth subsequent token candidates in parallel in the same forward pass. That many candidate tokens across each token position, from first to fourth, can be assembled into many candidate sequences, from which the model chooses the longest accepted candidate. In the case shown in Figure 7-2, one iteration of the model’s forward pass generates three tokens (It is difficult). Without speculative decoding, this pass generates only one token (It).

EAGLE does self-drafting differently than Medusa. It uses a small auxiliary module which is trained to predict the target model’s future internal hidden states (features) instead of tokens. This results in more stable and accurate predictions. Later on, EAGLE-2 introduced a dynamic draft tree to adapt the length of its speculation based on how predictable the current text is, and EAGLE-3 further enhanced accuracy by fusing features from multiple LLM layers to directly predict draft tokens. As we write this in fall 2025, EAGLE is one of the highest-performing self-drafting techniques. To perform well, however, it requires extra training and tuning.

N-gram

Instead of having a small model or module speculate about tokens, there’s a third option: a simple method called an n-gram. An n-gram basically chooses a sequence of n adjacent tokens from the earlier part of the request, saves them as a n-gram table, and performs matching to propose the next token(s). For example, say you give it the following:

	An hour ago, a quick brown fox ran away and now, that quick brown

Now you need to predict the next token. With n-gram, it captures this sequence in a cached n-gram table to propose that the next token will be fox.There’s no need for extra complex prediction. The original model takes fox and verifies it.

The n-gram method without an auxiliary model has very low computational overhead and can be very powerful, particularly in two scenarios. The first is when you are generating structured outputs such as JSON or SQL files, where the generation follows highly deterministic and repetitive patterns. The second is when asking the LLM to improve a piece of writing or fill a preset template–situations where LLM can quickly reuse many n-gram sequences from the original prompt, as shown in the following example.

Here is the original input prompt:

Write a brief confirmation email to a customer whose name is Iris for a Tuesday 2pm onboarding call. Keep it professional and friendly. Use the following email template:
Subject: <one line>
Greeting: Hi <name>,
Body: Confirm the onboarding call date and time; mention sending a calendar invite with a video link; offer to reschedule if needed.
Closing: Best regards

Here is the model’s generated response:

Subject: Onboarding Call Confirmation - Tuesday 2:00 PM

Hi Iris,

This is a quick note to confirm our onboarding call on Tuesday at 2:00 PM Pacific. I'll send a calendar invite with the video link shortly. If anything changes, just reply here and we'll find a time that works.

Best regards,
Andy

In this case, we are providing a lot of predefined structure in the prompt, which will help to achieve a higher acceptance rate. That, in turn, allows us to push for higher K and achieves lower latency.

Because n-grams incur minimal overhead, even if the acceptance rate is low, they are still a very viable option. We recommend making this the first speculative method you try before the other draft models.

Tune the number of draft tokens

The other crucial knob for speculative decoding is K, the maximum number of draft tokens. When K is large, it gives a very high ceiling for speedups. However, if the acceptance rate is low, it’s less effective, since many draft tokens are discarded and wasted. Conversely, if K is small, there are fewer tokens to verify and reject, so you get more predictable performance–but could be underutilizing speculative decoding and not getting its full benefit. In general, K values are something to tune on. They depend highly on the use case, the target models, the drafting method’s overhead, and the performance SLOs you want to achieve. For example, because the n-gram method has very low overhead, it can tolerate a lower acceptance rate and still provide a net gain in performance.

A smaller K, such as 2 or 4, can form a decent and stable baseline from which you can increment. In many cases, the best K is somewhere between 4 to 8. For very predictable generations, such as structured outputs or AI-agent function calling, K can be quite a lot larger, even pushing toward 16 or 32.

In many LLM serving frameworks, you can observe the acceptance rate for each position. For example, if K = 6, you might get [0.9, 0.7, 0.6, 0.5, 0.10, 0.02]. This means the first token has an 80% acceptance rate, while the fifth and sixth tokens have rates lower than 10%. In this case, you might want to consider lowering your K, since the last two speculations are not yielding much gain.

Limitations

Speculative decoding only helps with the Decode stage. Furthermore, the way it leverages a small draft model could result in extra computation when the acceptance rate is low. And if your scenario has a long input context in the Prefill stage, meaning that the model is already bottlenecked on compute FLOPS, speculative decoding might not work.

Likewise, processing large batches can shift the bottleneck from memory-bound to compute-bound, as we discussed extensively in Chapter 6. Even if your scenario does not have a very long context and generates a lot of tokens, efficient batching can still move the bottleneck from memory-bound to compute-bound. If so, speculative decoding can still help with latency, but the extra computation will increase overall throughput.

Productizing speculative decoding can be quite tricky. Not only is performance not-guaranteed, but running two models on a shared GPU efficiently and with high resource utilization is not a simple task. This is why speculative decoding has recently begun to focus more on self-drafting and n-gram techniques than on utilizing a small model. Also, a static K might not be able to capture the changing demand.

In short, the best scenario to utilize speculative decoding is when your workload is latency-sensitive and you are okay with trading some throughput and/or TTFT for better ITL and overall latency. When the Prefill:Decode token ratio is low and real-world batches are small, causing a GPU-memory-bound bottleneck, you can take advantage of speculative decoding’s blazing speed to meaningfully improve the customer experience.

Hands-On Speculative Decoding

Let’s try some speculative decoding examples in the Google Colab notebook using both n-gram and Eagle3:

# vanilla vllm
nohup vllm serve Qwen/Qwen3-32B \
 --disable-log-requests \
 --max-model-len 2048 \
 --gpu-memory-utilization 0.95 \
 > vllm.log 2>&1 &

# n-gram
nohup vllm serve Qwen/Qwen3-32B \
 --speculative-config '{
   "method": "ngram",
   "num_speculative_tokens": 6,
   "prompt_lookup_min": 4,
   "prompt_lookup_max": 6
}' \
 --disable-log-requests \
 --max-model-len 2048 \
 --gpu-memory-utilization 0.95 \
 > vllm.log 2>&1 &

# improved n-gram
nohup vllm serve Qwen/Qwen3-32B \
 --speculative-config '{
   "method": "ngram",
   "num_speculative_tokens": 4,
   "prompt_lookup_min": 2,
   "prompt_lookup_max": 128
}' \
 --disable-log-requests \
 --max-model-len 2048 \
 --gpu-memory-utilization 0.95 \
 > vllm.log 2>&1 &

# eagle3
nohup vllm serve Qwen/Qwen3-32B \
 --speculative-config '{
   "method": "eagle3",
   "model": "RedHatAI/Qwen3-32B-speculator.eagle3",
   "num_speculative_tokens": 3
 }' \
 --disable-log-requests \
 --max-model-len 2048 \
 --gpu-memory-utilization 0.95 \
 > vllm.log 2>&1 &

Running benchmarks

Inside the notebook, we run two small benchmarks: one with concurrency of 1 to simulate performance under low RPS, and the other with concurrency of 16 to simulate a high-load, high-RPS scenario over the four setups.

Performance analysis

Start with the concurrency = 1 results on the left side of Figure 7-3, where you can observe that n-gram slightly improves the token throughput. The improved n-gram provides a gain ~16% over vanilla vLLM. In comparison, Eagle3 achieves close to double vanilla vLLM’s throughput (56.5 tokens/s versus 28.9 tokens/s).

Figure 7-3. Total token throughput comparison between four different runs with concurrency of 1 and 16.

Moving to concurrency = 16 on the right of Figure 7-3, now both n-gram runs perform worse than vanilla vLLM. This is due to the overhead incurred with proposing and verifying new tokens. This causes extra CPU and GPU compute load that becomes even more substantial at high concurrency settings. Similarly, though Eagle3 still significantly outpaces vanilla vLLM, the gain is smaller.

Figure 7-4. TTFT and ITL trends across four different runs with concurrency = 16

To help you understand this further, Figure 7-4 shows the TTFT and ITL across four runs under a high RPS scenario (concurrency = 16). As expected, both n-gram variants and Eagle3 perform better than vanilla vLLM. However, Eagle3’s TTFT is significantly longer than that of vanilla vLLM, due to the overhead of the extra speculative heads and modules. Thus, the tradeoff is between prioritizing TTFT or ITL, based on your SLA or SLO.

In real-world settings, the setup will be highly dependent on the dataset, and you will need to perform multiple runs with understanding of the acceptance rate across K to select the decode method and the best value of K.

Muti-GPU and Multinode Inferencing

Modern LLMs often exceed the memory footprint of a single GPU, and serving them with low latency and at scale requires far beyond what one GPU can sustain. To overcome these limitations, inference systems turn to distributed strategies that harness multiple GPUs across one or more nodes. In this section, we shift our focus to distributed serving and discuss four parallelism techniques:

Data Parallelism (DP)

For scaling throughput; works by replicating model instances across GPUs and nodes

Tensor Parallelism (TP)

For fitting large models; works by splitting layers and large matrix operations to achieve lower latency

Pipeline Parallelism (PP)

For fitting large models; works by partitioning layers into stages across different GPUs and nodes

Expert Parallelism (EP)

For mixture-of-experts (MoE) models; works by distributing experts across GPUs

The next sections look more closely at each technique.

Data Parallelism

In software engineering, when the request count becomes too high, we horizontally scale the number of instances or replicas first. In the ML world, the preferred technique is very similar: we replicate the number of model instances so that each model instance can process a percentage of the total traffic. Both techniques are called data parallelism. DP is essential to scaling to handle large numbers of requests and ensure high availability. It is applied everywhere when handling production traffic.

Figure 7-5. Data parallelism applied to three model instances

Figure 7-5 shows a simple example of DP, where a load balancer/router accepts nine requests and spreads the load to three different model instances evenly. Each model instance gets three requests, so none of them is overloaded. If one of the model instances goes down, the router can stop routing requests to it and send them to the other two, as shown in Figure 7-6. (This is a simple example; in practice, the routing strategy can be much more complex than that.)

Figure 7-6. In data parallelism, when one model instance is down, the router routes requests to the two other running model instances

Common routing strategies used in DP include:

Round robin

The simplest method. Requests are sent to each model instance in a fixed, sequential order. While easy to implement, it can lead to uneven loads because some requests may take much longer to process than others, leaving some instances idle while others are overloaded.

Fewest connections

This strategy routes new requests to the instance with the fewest active connections. It’s an improvement over round-robin since it accounts for the current workload, but it doesn’t consider the complexity or length of the requests already being processed.

Latency-based routing

More sophisticated load balancers continuously monitor the real-time latency of each instance. Requests are routed to the fastest-responding instance. This method helps to reduce tail latency by dynamically adapting to performance fluctuations.

Cache-aware routing

In Chapter 6, we introduced prefix caching, a powerful technique to reduce TTFT by reusing existing KV caches across requests when prefixes are the same. As you scale to multiple model instances, the key to achieving a high prefix-caching hit rate is to route requests with the same prefixes to the model instances that already hold those caches (seeFigure 6-23 for a refresher).

Because LLM-serving workloads vary greatly from one request to another, and the underlying hardware for serving LLMs is extremely expensive, the value of an efficient routing mechanism is significantly higher than in traditional web applications. LLM routing has evolved beyond simple rule-based matching into a complex system that leverages real-time signals, such as:

KV cache locality

KV cache hit percentage (the ratio of number of tokens of KV cache hit to total prefix length)

Model instance’s KV cache space usage percentage

Raw input length

Number of requests queued for the model instance

Number of pending requests for the model instance

Number of tokens already generated for the requests in progress

Request SLA/SLO latency budget (the maximum acceptable response time to meet service guarantees).

To learn more about routing,see KV router by Nvidia Dynamo and the llm-d router documentation. These routers also support additional optimizations, such as Prefill/Decode worker routing and KV offloading, which we will cover in later sections of this chapter.

Tensor Parallelism and Pipeline Parallelism

When the data in a database grows too large for a single machine, we perform sharding to split the database into multiple shards, or pieces, across multiple machines. Similarly, when you need to perform inference with very LLMs, a single GPU often won’t be enough to load and serve the model fast enough to meet your SLA. This is where multi-GPU and multinode inferencing comes in. This technique involves distributing a single large model across multiple GPUs, which can be in a single machine (multi-GPU) or sharded across several interconnected machines (multinode).

Tensor parallelism (TP) and pipeline parallelism (PP) are two common methods for sharding a model across GPUs. TP shards the model along its width, while PP shards it along its depth, as shown in Figure 7-7. More specifically, TP splits the individual layers of a model horizontally across multiple GPUs. Each GPU holds a slice of a layer’s weight tensor and performs a partial computation. The results are then combined to produce the full layer output. Pipeline parallelism, in contrast, splits the model’s layers vertically across multiple GPUs rather than splitting any single layer into multiple pieces. Each GPU is responsible for a contiguous block of layers. Data flows sequentially through these GPUs, like an assembly line: one GPU computes the first few layers, then passes the intermediate results to the next GPU, and so on, until the final output is produced.

Figure 7-7. Serving a four-layer model with a single GPU (left), with two GPUs with TP = 2 (center), and with two GPUs with PP = 2 (right)

Figure 7-8 compares the two approaches. With TP set to 2, each layer of the model is split in half parts, with each half running on one of the two GPUs. With PP As for pipeline parallelism, instead of splitting each layer, different layers to to different GPUs. In the diagram, with four layers and PP is set to 2, layers 1 and 2 are located in one GPU and layers 3 and 4 in the other.

Figure 7-8. Cross-GPU communication in tensor parallelism (left) and pipeline parallelism (right)

In Figure 7-8, diving deeper, tensor parallelism needs inter-GPU (between GPU 1 and GPU 2) communication for every single layer. Pipeline parallelism only needs inter-GPU communication once it finishes layer 2 and moves on to layer 3.

Pipeline parallelism minimizes inter-GPU communication in processing, which is a good thing. Does this mean it’s better than tensor parallelism? It’s not quite that simple. Like a factory assembly line, PP can form bottlenecks that delay progress. Inside the LLM, workloads often vary across layers. When one stage becomes slow, other subsequent stages may be waiting, and devices become underutilized. This is called a bubble. For example, if layer 2 is very slow, layers 3 and 4 can be waiting for work. So GPU 1 could have very high GPU utilization while GPU 2 is idling most of the time.

Figure 7-9. Decision flowchart on when to use Tensor Parallelism vs Pipeline Parallelism

Now, let’s go through the decision flowchart in Figure 7-9 to better understand when to use these techniques and balance the trade-offs between TP and PP.

Starting from the top of the figure, the very first decision is whether you really need multi-GPU serving. You learned in chapter 5 that GPU memory bandwidth is much faster than GPU interconnect. For example, H100 can transfer about 3TB/s of data from GPU memory to the compute core; even with NVLink, the bandwidth is only ~900GB/s across GPUs, which is around three times slower. This wide performance gap in data transfer speed results in several considerations. First, whether you have high-speed interconnects like NVLink or not, if your workload can stay within one GPU, it is still better to keep it in one GPU and avoid communication overhead. Second, adding GPUs does not scale linearly: for example, adding one GPU to the existing one doubles your GPU memory space and GPU FLOPS, but does not double performance because of the interconnection bottleneck. Third, in many cases, getting a higher-end GPU with a bigger memory so that you can host the model in a single GPU is better than getting multiple lower-end GPUs with parallelism.

Another key consideration is that if you are serving a model that won’t fit in your GPU memory, don’t forget that you can use quantization, as we discussed in Chapter 6. For example, FP8 quantization cuts model size in half compared to FP16 or BF16, and W4A16 quantization can bring it down to a quarter of the original size, if accuracy allows.

Now suppose that, even with quantization, your model still does not fit into a single GPU’s memory, or that single GPU still isn’t powerful enough to meet latency SLA. The next step would be to consider getting a node with multiple interconnected GPUs. For example, one AWS P5.48xlarge node has eight H100 GPUs, interconnected with NVLink.

With one node containing multiple GPUs, you can start applying TP or PP so that you can run your model run across multiple GPUs. You’ve seen that PP minimizes the GPU interconnect load, but can cause pipeline bubbles. TP does not have this issue, but requires a lot of inter-GPU data transfer. Because of that, TP is the ideal choice if your GPUs within the node are connected with NVLink (for example, you’re running a big LLM on eight GPUs on P5.48xlarge–a pretty common setup). However, If you don’t have NVLink, current PCIe transfer speed is usually too slow for TP to work well. In that case, you consider switching to a different instance or trying PP– but, in general, this scenario is less common and not really ideal.

Figure 7-10. Serving a four-layer model with TP=8 and PP=2 on two nodes each, with eight GPUs

Finally, we move to the LLMs that require hyperscale setup, where one node with eight high-end GPUs still won’t be enough. In that case, pull in multinode serving. Since internode connection is even slower than intranode, this is where pipeline parallelism shines: it minimizes the required data transfer across nodes. Figure 7-1 shows an example of sharding a model both on width (TP = 8) and height (PP = 2).

This snippet is a simple example of enabling TP and PP within a single eight-GPU node:

hf_model_id = "Qwen/Qwen2.5-7B-Instruct"
!vllm serve \
     --tensor-parallel-size 4 \
     --pipeline-parallel-size 2

Running vLLM across nodes is a bit more complicated and less used because of the network bottleneck. If your workload requires multinode inference, you can use Ray with vLLM to achieve that.

Expert Parallelism

Modern LLMs often go beyond splitting a single dense model, using mixture-of-experts (MoE) layers designed to dynamically route tokens to different experts. Models such as Mixtral 8x7B, DeepSeek-V3 and GPT-OSS leverage MoE to reduce serving latency and cut serving costs while still delivering the performance characteristics of much larger, denser models. Figure 7-11 is a very simple illustration of how MoE works. To learn more on MoE, we recommend https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts.

Figure 7-11. Simple illustration of MoE with a router that determines which expert to be used. Image source: Maarten Grootendorst

While MoE architecture reduces serving-time computation by activating only a few experts per token, it introduces a new challenge: the total parameters can become too large to fit on a single GPU.Ideally, experts can be distributed efficiently across hardware. This is where expert parallelism (EP) comes in: it partitions the experts across multiple GPUs so that routing and execution can scale seamlessly, as shown in Figure 7-12. In practice, EP complements TP and PP to achieve the best performance.

Figure 7-12. Simple illustration of EP=2 for a four-expert model. (Other, non-MoE layers are not shown.)

Prefill-Decode Disaggregation

Tensor parallelism and pipeline parallelism mainly address how to fit and accelerate a single large model with multiple GPUs. However, they don’t solve inefficiencies that come from the two-phase nature of inference itself. Recall that the Prefill stage processes the entire input sequence in parallel and is compute-intensive, whereas the Decode stage generates tokens sequentially and is dominated by memory bandwidth. In all of our current setups, these two stages are scheduled on the same GPU or GPUs. This colocation usually means that the GPU needs to process one phase, either Prefill or Decode, while the other waits. In Chapter 6, you learned that the chunked prefill technique can mitigate this issue by splitting long input prompts into smaller chunks, thus unblocking the Decode phase from waiting for long prefill workloads to complete. However, processing Prefill and Decode at the same time on the same set of GPUs still presents a fundamental challenge: the two stages have vastly different resource-utilization profiles. This can cause high interference and makes it difficult to optimize for both simultaneously.

Thus, separating the Prefill and Decode stages into different GPUs has become a popular method for large models and large-scale serving systems. By physically separating the compute-heavy Prefill phase from the memory-and-bandwidth-heavy Decode phase, Prefill–Decode (PD) disaggregation provides several important benefits:

Optimize TTFT and ITL independently

On the user-experience side, PD disaggregation allows you to independently optimize TTFT or ITL by scaling the resources for Prefill and Decode asymmetrically, as shown in Figure 7-13. Depending on the use case, for example, the input/output token ratio can vary a lot. For an input-heavy, interactive workload, TTFT is key. PD disaggregation makes it easier to add more resources dedicated to the Prefill stage to complete long input prompts much faster, reducing wait time for the first token generation (TTFT). You can do the same for Decode later on, if ITL needs improvement. Moreover, because there is no interference due to having Prefill and Decode on the same compute resource, ITL can be stabler and more predictable.

Figure 7-13. Different use cases can have very different input (Prefill) and output (Decode) ratios, leading to different TTFT/ITL tradeoffs and requirements.

Separating the Prefill and Decode workloads

On the GPU-utilization side, PD disaggregation lets you completely separate the Prefill and Decode workloads. The high-compute Prefill stage does not sit idle waiting for the memory-bound Decode stage to finish, and vice versa. This separation also lets Prefill and Decode have different batch sizes.

For Prefill, depending on the input size, you can adjust the batch size to produce large enough matrix multiplications to saturate GPU tensor core utilization. For Decode, you can tune the batch size to be large enough to ensure that the workload is no longer memory-bandwidth-bound while balancing the ITL latency requirements. Similarly, you can apply a heterogeneous parallel strategy: for example, you might apply TP and PP differently to Prefill or Decode to further improve or fine-tune the performance of one stage without affecting the other.

Opening up more hardware options

On the hardware selection side, PD disaggregation also means more flexibility. Prefill can use more compute-optimized GPUs for more FLOPS over cost, while Decode can be scheduled on memory-optimized GPUs, with high memory bandwidth over cost. For example, the H200 GPU and H100 GPU have the same FLOPS specifications, but the more expensive H200 has a bigger GPU memory (141GB over 80GB) and GPU memory bandwidth (4.8TB/s over 3.35TB/s). In this case, the H100 could be more cost-effective for the Prefill stage, while the H200 would be better for Decode. If you don’t need that much performance for Decode, you could use the L40S instead, which has decent GPU memory for a much lower cost. In this example setup, using the H100 for Prefill and the L40S for Decode, you can achieve an optimal TTFT and satisfactory ITL while saving a significant amount of money.

Enabling independent and asymmetric scaling

PD disaggregation offers fine-grained control of the Prefill and Decode stages, which means their scaling can be more elastic to deal with fluctuating demand. Prefill, as the entrypoint of the requests with varying input length, can be more burst-driven and can be paired with a more aggressive scaling strategy. Decode generates tokens over a more prolonged period and is more predictable and stable.

Overall Architecture

Figure 7-14. Prefill-Decode disaggregation architecture from DistServe. (Image source: Zhong et al., 2024)

Now let’s look at how it’s done. Figure 7-14 shows the original DistServe runtime system architecture. All incoming requests arrive at the controller for routing. Requests are first sent to Prefill instances for the Prefill stage. The Prefill instances produce the KV cache, which is immediately transferred to the Decode instances, which then generate tokens to complete the requests.

Figure 7-15. Side-by-side comparison of traditional aggregated serving andPD disaggregated serving (Image source: NVIDIA)

Figure 7-15 shows that in traditional aggregated serving, GPU instances are tasked to do both Prefill and Decode, consistently balancing the compute-bound and memory-bound workloads. In disaggregated serving, however, GPU instances are split into two instance groups, with dedicated Prefill instances to maximize computation and Decode instances to tackle memory-bandwidth issues.

KV Cache Transfer

One important step we have not yet covered, but which you might have noticed in both Figures 7-15 and 7-16, is the KV cache transfer. Efficiently transferring the KV cache from the Prefill instances to the Decode instances is critical to the success of disaggregated serving. The performance gain from PD disaggregation needs to outweigh the overhead of the KV cache transfer.

As we discussed in Chapter 5, GPU memory is scarce and its most significant usage, apart from model weight, is for KV cache, which is essential for fast decoding when you trade memory space for less computation. Despite technological advances and compression and quantizations techniques to reduce its size (see Chapter 6), the KV cache can be very large. For example, a common 8-billion-parameter model with 1,024 input tokens can have a KV cache size of around 0.1 to 0.15GB. This might not sound like a lot, but consider that a GPU is not processing one request at a time. As continuous batching pushes the throughput higher, one GPU can process tens of requests or more per second. Moreover, in real-world applications, input length can easily be multiples of 1,024. The KV cache size to be transferred scales mostly linearly with input length. Thus, 10x means KV cache size per request becomes 1 to 1.5GB. Now multiply that by a throughput of, say, 16 requests per second, and you need to be able to transfer close to 25GB of KV cache per second.

with NVLink, Intra-node communication bandwidth performance can be as high as 900GB/s or more, whereas inter-node, even with InfiniBand, is around 50GB/s to 100GB/s. (See Table 5-3 for a comparison of intra-node and inter-node communication bandwidth.) Without RDMA technology like InfiniBand, inter-node connection needs to go through PCIe, which is usually around 10GB/s. Thus, how to allocate the GPUs for Prefill and Decode based on your hardware setup becomes important.

This example’s 25GB KV cache transfer is well within the capabilities of intra-node NVLink (900GB/s) and works okay for inter-node InfiniBand (50-100GB/s), but exceeds the practical limits of inter-node through PCIe when RDMA is not available. In such cases, it is crucial that the Prefill and Decoding instances for a given request are handled by GPUs within the same server node, so communication is done via NVLink, as shown on the right side of Figure 7-16. On the contrary, the allocation on the left forces the massive KV Cache to be transferred over the much slower inter-node network. This can create a severe communication bottleneck, effectively saturating the network fabric and throttling the overall system throughput.

Figure 7-16. Two different GPU allocation comparisons, with inter-node communication on the left and intra-node communication on the right.

Restricting intra-node placement of Prefill instances and Decode instances is not ideal. It forces you to use the same GPU hardware for both Prefill and Decode. It also limits model parallelism in many cases. For example, let’s say one model instance is big enough that you want to use tensor parallelism with eight GPUs for prefill. In this case, inter-node communication would no longer be optional and would provide much more flexibility for scaling, especially when the Prefill and Decode workloads change. To combat the overhead of KV cache transfer in an inter-node setup, using RDMA across nodes is still a must, but there are several optimizations you can do to further reduce the overhead. For example:

Instead of transferring the entire KV cache only when completely finished, you can chunk it to be transferred in smaller blocks, much like how streaming video works. You can tune the size of the blocks for the best performance.

The transfers use asynchronous, nonblocking operations, so compute and communication can overlap as much as possible. When the Prefill GPU is still performing the computation, the transfer is happening behind the scenes so the Decode GPU can start its work sooner. This effectively hides the data transfer behind the stages of compute.

Because the KV cache is per layer and is completely layer-local, it can also be transferred layer by layer. As the Prefill stages computes the subsequent layers, the KV cache of the completed layers will be transferred, so the Decode stage can start on them.

You can reduce or compress the KV cache’s size.

Figure 7-17. Simplified illustration of how the KV cache transfer can be hidden and overlap with compute

With all these optimizations on KV cache transfer and improved scheduling, you can reduce the overhead of PD disaggregation to less than 1% of the overall per-request latency (Wang et al. 2025). Figure 7-17 shows a very simplified illustration of how KV cache transfer is hidden and overlaps with GPU compute so that, when managed correctly, it causes minimal overhead.

When to Use

Figure 7-18. Decision flowchart for when to go with aggregated serving vs. PD disaggregated serving

With all the benefits of PD disaggregation, you might be wondering if you should use this technique by default. The answer is no. Figure 7-18 shows a detailed decision flowchart to help you decide which kind of serving to use based on your stage. But, simply put, use PD disaggregation only for large models under heavy load with advanced TTFT and ITL tuning needs, and stick with simpler setups for smaller models.

Advanced KV Caching

Up to this point, we have explored techniques that make LLM serving faster and more efficient: speculative decoding to shorten token-generation time, multiple parallelism strategies to utilize multiple GPUs for large models, and PD disaggregation to better balance the distinct demands of different stages of generation.

In real-world serving, we are always pushing the boundaries further, especially given modern applications; increasing demand for models that can reason over vast amounts of context to provide tailored, customized responses. For example, coding copilots must keep track of thousands of lines of source code so the model can generate changes in the right context. Conversational agents must sustain long interaction histories while drawing on tenant-specific knowledge bases to deliver consistent answers. Enterprise platforms must handle large document collections or transaction histories to extract insights, detect anomalies, or personalize decisions.

Long-Context Serving

Retrieval-augmented generation (RAG) is a very common strategy that retrieves the most relevant documents or records at query time and feeds them into the prompt, as you learned in Chapter 4. This approach keeps the context length and TTFT in the Prefill stage manageable, while grounding responses in fresh, tenant-specific knowledge. RAG has proven invaluable in domains like customer support, enterprise search, and analytics, where accuracy heavily depends on dynamically retrieving the right slices of tenant-specific documents, records, or interaction histories. Cache-augmented generation (CAG, also discussed in Chapter 4) tries to cache most or even all of the relevant context as KV cache and treats it as reusable knowledge across different requests.

Prefix caching could be considered a naive version of CAG without advanced optimizations: you could treat all the relevant knowledge as static prefixes of the input, and only append the dynamic prompts as suffixes to form the full input to send to the LLM. LLMs can leverage prefix caching to reuse the KV cache, achieving very low TTFT. The idea has been there for quite some time, hasn’t been implemented and adopted as fast as RAG has, only picking up more recently.

There are shifts happening on two fronts that make long-context CAG-based serving more plausible. First, we’re now seeing the rise of long-context models, with windows stretching from 100k to 1M tokens. This makes it possible to load most or all of the tenant’s entire context directly into the model. In other words, rather than repeatedly relying on retrieval to inject information into each request, these models can directly attend to entire histories, large document sets, or transaction logs in a single extended context window. As these long-context LLMs become more and more capable, they are solving the “lost in the middle” problem, where LLMs are positionally biased to emphasize information at the beginning and ending over that in the middle.

On the engineering side, GPU memory is still very limited, despite newer chips increasing memory capacity. However, with efficient KV cache management techniques like offloading to CPU memory, SSD, or even remote storage, along with cross-replica routing, it is now very plausible to have an LLM cache large amounts of knowledge as KV cache to be reused when needed.

This shift reduces the complexity of managing retrieval pipelines, makes reasoning over large amounts of input smoother, and in many cases produces better responses. Without real-time retrieval, the long-context CAG approach can significantly reduce redundant computation, resulting in faster TTFTs than RAG achieves.

Cost and Latency Calculations

Figure 7-19. Example input-token comparison between RAG serving (top) and long-context CAG serving (bottom)

Figure 7-19 shows compares input tokens from RAG and long-context CAG serving. In this example, both have a system prompt with 500 tokens and can be easily cached in the KVcache to avoid recomputation across requests.

For RAG, depending on the user prompt, the model will dynamically generate the relevant knowledge chunks: in this case, 10 different chunks per 500 tokens. Because they vary by the user prompt, the KV cache hit rate will be very low. Adding the user prompt, LLM needs to run Prefill on:

0 + 500 × 10 + 500 = 5000 tokens

With long-context CAG, both the system prompt and the static long context can be retrieved from the KV cache, so only the user prompt is dynamic. What we run Prefill on is thus, essentially, 0 + 0 + 500, or 500 tokens.

Because CAG only takes in 1/10th of the input for Prefill, it gains correspondingly in TTFT time. So if, with RAG, your TTFT is 5 seconds, long-context CAG could shave it down to around 0.5 seconds. In many production use cases, that’s a night-and-day improvement. This doesn’t even account for RAG embedding time and vector-search time prior to LLM generation, so the real-world TTFT savings would be even bigger.

However, you do need to account for the cost for caching long context, which will depend on whether the model is self-serve or uses an external vendor API. For example, Table 7-1 compares the costs for regular input versus cached input as we write this in late 2025. As you can see, the cached input is usually 10 to 25% of the regular input.

Table 7-1. Cost comparison of regular input and cached input across external vendors, in $ per 1M tokens

GPT-5
Gemini 2.5 Pro
Claude Sonnet 4

Regular Input

1.25

1.25

3

Cached Input

0.125

0.31 (+ 4.5 storage/hr)

0.3 (+ 3.25 first write)

Let’s include the GPT-5 costs in our example calculation. For RAG, that would be:

(500 × 0.125 + 500 × 10 × 1.25 + 500 × 1.25) / 10⁶ = $0.007 per request

For long-context CAG:

(500 × 0.125 + 100,000 × 0.125 + 500 × 1.25) / 10⁶ = $0.013 per request

Long-context CAG is almost double the cost of RAG here! This is because we have 100,000 tokens of cached inputs rather than 5,000 tokens of regular input. This calculation doesn’t account for the RAG costs of offline indexing, storage and online retrieval, either–but those usually cost much less than LLM calls, so RAG will still be cheaper even though TTFT is a lot slower.

Self-Hosting LLMs

All of the calculations above assume that you are using a third-party API. Now we’ll look at how to optimize a self-hosted model to make long-context CAG work faster and cheaper.

KV Cache Offloading

At this point, the KV cache is no longer just an implementation detail hidden inside the model—it needs to become a first-class citizen in the serving stack. It must be stored, scheduled, moved across devices, and evicted with the same care as model weights or input data. Systems such as LMCache are designed around exactly these aspects, elevating KV management into a core responsibility of large-scale LLM serving systems (see Figure 7-20).

Figure 7-20. Tiered storage hierarchy for KV Cache storage, with example capacities based on an AWS p5.48xlarge instance

The very first key feature that LMCache provides (as do a few other similar frameworks) is efficient tiered KV cache offloading. Figure 7-21 shows that offloading GPU memory to CPU memory can provide around triple the KV cache space. With SSD, you can get as much as 50 times more KV Cache space! This means you can cache 50 times more long-context documents for many different uses or cache multiple tenants’ documents within one instance. And all that extra space is literally free if you are already self-serving on the same instance. Adding distributed network storage, such as Redis or even S3, can be an option for more storage, should you need it.

Figure 7-21. Long-context serving without KV cache offloading (left) and with KV cache offloading (right)

Figure 7-21 shows how KV Cache offloading can help save costs over vanilla prefix caching. On the left, with prefix caching, no single model instance has enough GPU memory to hold more than one long-context KV cache. This forces us to create multiple model replicas in combination with a routing layer to ensure KV cache hits. On the right, with KV cache offloading, in this case only using CPU memory to triple the space, one model instance can keep all four long-context KV caches and, depending on the request prefixes, swap between GPU and CPU memory efficiently. If requests for one long context or one tenant cannot saturate a single model instance, this translates to fourfold cost savings and a million levels per year.

KV Cache Compression

To further improve f KV cache transfer performance and efficiency, especially when load is high and transfer bandwidth is limited, compressing the KV cache compression can really help. You can do this using traditional KV cache quantization or the more advanced CacheGen by LMCache, which encodes the KV cache into more compact bitstream representations based on its distribution, to ensure high quality and reduce size significantly. The smaller KV cache can be transferred much more quickly, with fewer network delays, and without too much GPU computation overhead during decompression. Combined with offloading, this also means more space in GPU memory to handle higher batches, increasing throughput and caching more context per model instance.

KV Cache Blending

Even all these KV cache management techniques don’t mean that RAG is obsolete. Real-world enterprise data can be huge, highly dynamic, and often unbounded in scope. RAG still has the advantage in accessing a larger knowledge base, retrieving fresher information without data staleness, and potentially being easier to explain and trace for citations.

Figure 7-22. Naively concatenating KV caches does not work, because self-attention requires the later tokens to attend to the entire prior history. (Image source: Yao et al., 2025)

One novel, RAG-specific KV cache management idea is CacheBlend, supported by LMCache. CacheBlend lets you merge KV caches of different RAG chunks together, to reduce the amount of recomputation of all RAG chunks from scratch. Recall that LLMs can only match prefixes; different RAG chunks, or the same RAG chunks in a different order, can cause the KV cache to miss and do a full recomputation. You might think this could be solved by simply concatenating different KV caches, but it’s not that straightforward. Self-attention doesn’t just attend within each cache independently–it also learns cross-token interactions between caches. A naïve concatenation would discard these relationships, effectively breaking the contextual dependencies that the model could have learned. Figure 7-22 provides an example.

To combat this issue and still leverage existing non-prefix KV caches, CacheBlend selectively recomputes a small subset of KV cache (the default is 15%, which can be tuned with blend_recompute_ratios) to preserve cross-token relationships. This avoids incurring the cost of full recomputation while achieving a very minimal drop in quality.

Hands-On LMCache

In this section, we’ll show you some examples of how to enable LMCache in practice.

In the Google Colab notebook that accompanies this book, you can find the code that hosts this Qwen/Qwen3-14B model in two different configurations: vLLM serving with LMCache enabled for local CPU-based offloading, and vanilla vLLM serving with default prefix caching.

Enable LMCache and CPU offloading

First, let’s look at an example command to enable LMCache with vLLM server. Here, we enabled local CPU memory for KV cache offloading and set the CPU memory space to be 150GB.

export LMCACHE_USE_EXPERIMENTAL=True
export LMCACHE_LOCAL_CPU=True
export LMCACHE_MAX_LOCAL_CPU_SIZE=150.0

# Uncomment to test SSD offloading
# export LMCACHE_CHUNK_SIZE=2048
# export LMCACHE_LOCAL_DISK="local_kv_cache/"
# export LMCACHE_MAX_LOCAL_DISK_SIZE=120.0

# Start the vllm server
nohup vllm serve \
  Qwen/Qwen3-14B \
  --disable-log-requests \
  --kv-transfer-config '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_both"}' \
  --max-model-len 40960 \
  --gpu-memory-utilization 0.9 \
  > vllm.log 2>&1 &

Run benchmarks

In the example notebook, we ran two different benchmarks to show the performance difference between vanilla vLLM and vLLM with LMCache enabled. First, we made 30 sequential requests, each with unique prefixes. Then, we called those 30 sequential requests in reverse order to show the difference in KV cache space limit between LMCache with CPU offload enabled and vanilla vLLM using GPU memory only. Because the KV cache eviction setting defaults to Least Recently Used (LRU) and KV cache space is limited, if we send new requests, we will not get the prefix cache hits anymore and latency will go up.

In the second benchmark, we use the vLLM bench CLI to submit 20 different long-context requests in three waves. First, we start with everything cold and send 20 requests sequentially, so neither the LMCache setup nor the vanilla vLLM has any KV cache. In the second wave, we send the same 20 requests again, sequentially. Last, we send the same 20 requests, but with 5 concurrent requests at a time, which helps us to analyze the performance change under higher load.

Performance analysis

Figure 7-23. Execution time for 30 sequential requests sent to vanilla vLLM and LMCache with CPU offloading enabled during cold start without KV cache hit

Figure 7-24. Execution time for 30 sequential requests sent to vanilla vLLM and LMCache with CPU offloading enabled and KV cache stored

As Figure 7-23 shows, when there are zero KV cache hits during the first run, LMCache is always slower than vanilla vLLM. This is likely due to the overhead of KV cache handling. This means that unless you expect your workload to have a lot of KV cache hits, vLLM without LMCache performs better and should be the default setup when prefixes are short or requests are very random.

However, during the second run, after warming up, as shown in Figure 7-24, for the first nine requests, which are all within GPU memory for both setups, they are neck and neck. . After the 10th request, vanilla vLLM starts to reperform Prefill computations, and latency shoots back up to the level of cold-processing latency. This happens because the older cache is all evicted while LMCache starts to retrieve the cache from CPU memory. The latencies when using CPU memory are not as good as when retrieving directly from GPU, but still much lower than performing full Prefill computations (1 second versus 6 to 7 seconds).

Figure 7-25. Comparing median TTFT of LMCache and vLLM, with sequential requests (left) and 5 concurrent requests (right).

Figure 7-26. Comparing throughput with LMCache and vLLM, with sequential requests (left) and 5 concurrent requests (right).

During the second benchmark (Figure 7-25), we use vllm benchmark cli to show how the performance changes under load. Because we are sending 20 different prefixes, vLLM can no longer store them all, making its median TTFT a lot higher than LMCache;s. The gap is even bigger once we put in more concurrent requests: the vanilla vLLM setup gets bottlenecked by the KV cache space, because the same GPU memory needs to handle the KV cache and activations of five concurrent requests, which further squeezes the number of KV caches that can be stored. This is even more apparent in the throughput metrics (Figure 7-26), where LMCache easily handles the increased load while vLLM is bottlenecked at around 4,700 tokens per second–around 16 times less than the throughput of the LMCache setup.

Summary

In this chapter, we explored advanced techniques for large-scale LLM serving and optimization. We began with speculative decoding, a powerful strategy to reduce latency by predicting multiple tokens in parallel and verifying them efficiently. We then examined different model parallelism techniques and combinations such as tensor parallelism (TP) and pipeline parallelism (PP), which enable scaling large models across multiple GPUs by splitting computations along dimensions or sequential layers. Building on that, we discussed Prefill-Decode disaggregation (PD), which decouples the Prefill and Decode stages so each can run on the most suitable hardwares and configurations, enabling finer-grained scaling and better hardware utilization. Finally, we examined advanced cache-management techniques such as offloading KV caches across memory and storage tiers (GPU, CPU, SSD) to balance speed and capacity, cache compression methods that shrink memory footprint with minimal accuracy loss, and novel KV cache-blending techniques to merge non-prefix KV caches for RAG usage.

Modern LLM serving has evolved far beyond running a single model instance efficiently. Today it is a hierarchical stack, as shown in Figure 7-27. At the lowest level lie the highly optimized CUDA and Triton kernels, which maximize GPU compute and memory utilization efficiency. Above that, the execution-engine layer handles token scheduling, batching, and streaming, and coordinating Prefill and Decode execution across GPUs. The cache-management layer maintains KV cache allocation, compression, and tiered offloading across GPU, CPU, SSD and distributed storage to balance latency and capacity. At the level above, the routing layer operates to direct each request to a Prefill or Decode worker and to the optimal model instance, based on real-time signals such as KV cache locality, queue depth, and latency budgets. Above it all, the orchestration layer governs global resource management and autoscaling, using aggregated utilization and latency metrics to optimize throughput and cost efficiency across the entire fleet.

Figure 7-27. NVIDIA’s Dynamo architecture (Image source: NVIDIA)

Together, all these layers form an integrated optimization stack to deliver state-of-the-art latency, throughput, and cost efficiency. Achieving production-grade performance requires tuning across all layers in concert, since bottlenecks at any single stage can erase gains elsewhere. For organizations operating LLMs at scale, such system-wide efficiency is not just a technical pursuit: it can be both a use-case enabler and a core economic differentiator.

This concludes our three-chapter focus on LLM optimization techniques. Across these chapters, we’ve examined why and where serving complexity arises, explored the fundamental techniques that address it, and then advanced to multilayered optimizations that unify compute, memory, and orchestration. Together, they provide the foundation for designing LLM systems that are scalable, performant, and cost-efficient in real-world production environments.

In the next chapter, we’ll look closely at the vLLM framework to show you how to implement a highly optimized LLM serving stack.

¹ To dive deeper into the mathematics behind speculative decoding and the modified distribution, refer to the original research papers: "Fast Inference from Transformers via Speculative Decoding" (Leviathan et al., 2022) and "Accelerating Large Language Model Decoding with Speculative Sampling" (Chen et al., 2023).

Chapter 8. LLM Serving Frameworks

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 8th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at sgrey@oreilly.com.

Over the previous chapters, we’ve explored the fundamentals of LLM serving—system design, service implementation, and practical optimization techniques. This chapter shifts to the foundation layer–the serving frameworks that implement and execute model inference with different optimization techniques under real production constraints.

We’ll discuss four widely adopted open-source serving frameworks you’re likely to encounter in the wild: vLLM, TensorRT-LLM, SGLang, and llama.cpp. Each has a distinct philosophy, hardware footprint, and battle-tested technology, and is backed by active communities and growing production usage.

Because it’s the most broadly applied framework, we’ll take a deep dive into vLLM—its architecture, initialization and model-execution process, request and token-level scheduling, and layered optimization strategy. Understanding vLLM’s internals will give you strong intuition for how LLM frameworks work in practice and make it easier to evaluate the trade-offs in other frameworks.

Next, we’ll cover the remaining frameworks with concise, decision-oriented overviews and short examples. We’ll close the chapter with the evaluation method we use to compare serving frameworks.

After reading this chapter, you’ll have a solid grasp of what LLM serving frameworks are, why we need them, how they work under the hood, and how to evaluate them for your use case.

In the next chapter, we’ll put the optimization techniques and serving framework into practice—tuning a LLM performance with the vLLM serving framework.

Why We Need Specialized LLM Serving Frameworks

Before the LLM era, there were already many general-purpose model-serving frameworks, such as TensorFlow Serving, TorchServe, and even generalized inference platforms like NVIDIA Triton. These frameworks and platforms were originally designed for deep-learning workloads like image recognition and structured data inference. Those workloads usually have short input sizes, fixed-shaped tensors, and predictable latency requirements, and their primary optimization is usually batch processing.

At this point in the book, it should be clear that serving LLMs is fundamentally different from serving traditional machine-learning models, such as image classifiers or recommendation models. LLM serving and optimization introduce a new set of challenges, including:

Autoregressive generation

LLMs generate outputs one token at a time. Unlike with image models, inference sessions may remain open for seconds or minutes.

Context length explosion

Models must handle input prompts of sizes ranging from a few tokens to hundreds of thousands or even a million tokens. KV cache memory management becomes a critical bottleneck.

Continuous batching

Input and output lengths vary wildly across requests. Static batching strategies underutilize GPUs.

Streaming requirements

Users expect time-to-first-token (TTFT) within hundreds of milliseconds, as well as continuous token streaming.

Resource utilization

GPUs are expensive. Wasting GPU FLOPs due to fragmentation or idle tokens is unacceptable at scale.

To meet these needs, a new class of frameworks has emerged–specialized LLM serving frameworks, such as vLLM, TensorRT-LLM, and SGLang. These frameworks introduce innovations such as paged KV caching, continuous batching, LLM-specific quantization, and speculative decoding (discussed in previous chapters) to address the unique challenges of LLMs.

With the help of these LLM serving frameworks, we can squeeze more throughput and lower latency from modern accelerators and deliver far better efficiency–which is why they have become the mainstream choice for LLM serving.

vLLM

vLLM is one of the most popular open-source frameworks for LLM serving. Its adoption has been rapid across both open-source communities and enterprise teams because it solves the pain points of running LLMs at scale: long prompts, high memory demands, and the need to serve many users simultaneously.

At its core, vLLM introduces paged KV caching and incorporates it with continuous batching, two innovations that dramatically improve throughput and reduce latency. It also supports features like quantization, speculative decoding, streaming responses, and multi-GPU and distributed execution. These capabilities suit it to a wide range of scenarios—like interactive chatbots and RAG systems, batch text generation, multitenant serving, and real-time applications.

Many people like vLLM because it “just works.” It’s battle-tested, it’s easy to integrate with open source models and self-fine-tuned models, it maximizes GPU efficiency without requiring deep system-level tuning, and it delivers predictable performance gains out of the box. This mix of usability and efficiency is why vLLM has become a go-to choice for LLM serving in production.

Another strength of vLLM lies in its architecture. Its clean, extensible design makes it easy to incorporate the latest academic advances and practical optimizations. Combined with its large, active community, vLLM is not only well suited to today’s serving challenges but is also positioned to adapt and thrive as future needs evolve.

In this section, we’ll take a deep dive into vLLM—starting with its high-level system design, then walking through the model initialization workflow and request-handling pipeline. From there, we’ll move into more advanced topics, such as how generation requests are prioritized and optimized at the framework level to support a wide range of LLMs.

vLLM’s Architecture

vLLM is optimized for a single model service setup, in which each instance initializes one model at startup. It offers two primary ways to use that model:

LLM class (in-process library)

A pure-Python local interface for offline inference—no separate server or web API required. This “library mode” is ideal when you want to plug vLLM directly into an existing service or batch workflow.

API server (OpenAI-compatible)

A standalone HTTP server for multiclient and production use and streaming, with easy integration via the Chat/Completions API.

Using the LLM class gives you tight control, low overhead, and easy integration with your own serving stacks, while the API Server provides a ready-to-use network endpoint for broader deployment. Let’s see the LLM class example first.

# initialize LLM model
llm = LLM(
  model="Qwen/Qwen3-7B-Instruct",
  trust_remote_code=True, # Qwen uses custom modeling code
  dtype="float16",  # Use float16 for GPU
  max_model_len=32768,
  gpu_memory_utilization=0.8,
)

# run model generation requests
outputs = llm.generate(prompts, sampling)

Start the vLLM API server with the vllm serve command.

# start vLLM api server
vllm serve Qwen/Qwen3-7B-Instruct \
  --trust-remote-code \
  --dtype bfloat16 \
  --max-model-len 32768 \
  --gpu-memory-utilization 0.8

# call the Qwen model via the OpenAI-compatible API
curl http://localhost:8000/v1/chat/completions \
 -H "Content-Type: application/json" \
 -d '{
   "model": "Qwen/Qwen3-7B-Instruct",
   "messages": [{"role":"user","content":"Explain KV cache in two sentences."}],
   "temperature": 0.7,
   "max_tokens": 256,
   "stream": true
 }'

Now let’s turn to vLLM’s major internal components, as shown in Figure 8-1.

Figure 8-1. The vLLM system architecture

LLMEngine and EngineCore

LLMEngine is the high-level interface and main entry point for vLLM’s inference system. It serves as the public API that users interact with, while internally coordinating all the underlying components.

LLMEngine integrates all components into a cohesive system with clear data flow. It handles both synchronous and asynchronous serving scenarios and manages the overall request lifecycle, such as orchestrating the request-processing pipeline and managing the request queue and configuration.

EngineCore is the central orchestrator of vLLM’s inference engine. It integrates the model executor, output processor and scheduler, and serves as the “inner loop” that coordinates all major components and manages the complete request-processing pipeline.

Scheduler

The Scheduler component acts as the “traffic controller” for the entire inference pipeline. It manages the compute resources and orchestrates token calculation across requests. Its primary responsibility is to allocate limited computational resources (such as GPU memory, KV cache blocks, and processing capacity) efficiently among competing requests while maximizing throughput and maintaining fair access.

From an optimization perspective, the Scheduler is responsible for systemwide, model-agnostic optimization strategies such as token-level scheduling, dynamic batching, prefix caching, and chunked prefill. In contrast, model-specific optimizations are applied within the ModelExecutor and GPUWorker, which we’ll discuss next.

Scheduler encapsulates its model-execution plan into a comprehensive data structure called SchedulerOutput, which serves as the “work order” from the scheduler to the ModelExecutor, containing all the information needed to execute a batch of requests. The ModelExecutor responds to run the model. The SchedulerOutput essentially tells the ModelExecutor: “Here’s a batch of requests to process. Here’s how many tokens each one should get, here are their input data and parameters, here are the memory blocks allocated for them, and here are any special processing requirements.”

The ModelExecutor uses the information in SchedulerOutput to prepare the actual GPU batch (with flattened input IDs, attention metadata, and KV cache blocks) and execute the model forward pass, then return the results back to the scheduler for the next iteration.

ModelExecutor, (GPU) Worker, and ModelRunner

Since vLLM hosts and runs each model in its own separate process or process group, it employs a layered architecture to handle the cross-process communication, orchestrate distributed worker groups, and handle various model forward-pass execution details. This architecture includes three components:

ModelExecutor orchestrates and manages multiple worker processes.

GPUWorker works as the worker interface that runs in each worker process and manages device/model lifecycle

GPUModelRunner actually executes and runs the neural network

This separation of concerns allows each component to focus on its specific model-execution responsibility while maintaining clean interfaces between them. Figure 8-2 illustrates this layered cross-process model-execution architecture.

Figure 8-2. The relationship between ModelExecutor, GPUWorker, and GPUModelRunner

With this architecture in mind, let’s now explore how these components come together in practice—from model initialization to request execution.

Model initialization workflow (with multi-process worker)

In this section, we’ll walk through the model initialization workflow in vLLM—covering how it brings up internal components, sets up the worker process group, loads the model from its source into GPU memory, and establishes the communication links.

Because vLLM supports a wide range of models and multiple execution styles—single-device serving, multi-device serving on one node, and even multi-node clusters—the initialization logic can feel complex. To make it easier to follow, we’ll focus on a practical example: initializing an LLM with a multi-process setup, the configuration most commonly used in production today. See Figure 8-3 for a visualization of this process.

Figure 8-3. The vLLM initialization workflow with a multi-process serving setup

Let’s walk through the four steps shown in Figure 8-3:

1. Initialize all components in the vLLM main process

When creating an LLM() instance, we pass in configurations that define the model-setup, resource-usage, and serving-optimization parameters. In the following example, we configure vLLM to host a Qwen model from Hugging Face using a multi-process group with four worker processes. We used a multi-process (MP) backend for single-node model deployment with multiple GPUs, but you can also choose the “ray” backend for cross-machine setup.

llm = LLM(
  model="Qwen/Qwen2.5-7B-Instruct",
  # specify 4 workers
  tensor_parallel_size=4,
  # use multi-process model executor
  distributed_executor_backend="mp"
            )

The LLM class initializes all major components within the vLLM main process and distributes the configuration to the appropriate modules—including the LLMEngine, Scheduler, KV Cache Manager, and MultiProcessExecutor.

2. Create a worker process group

During initialization, the MultiProcessExecutor spawns four worker processes for distributed model execution. It also sets up an rpc_broadcast_mq message queue to pass signals and commands to these workers.

3. Initialize the worker processes

Each worker process runs a GPUWorker, which is responsible for executing model inference and communicating with the MultiProcessExecutor. The GPUWorker sets up the CUDA device, establishes cross-process communication, and loads and initializes the model within the worker process. For example, the GPUWorker maintains a message queue (worker_response_mq) to send inference results back to Model Executor.

4. Prepare and load the model

The GPUModelRunner is responsible for locating the correct model implementation based on the configured model name. In our example, it looks up the Qwen model in vLLM’s internal model registry and selects the Qwen3NextForCausalLM implementation. It then calls the class’s __init__ function to load the model weights onto the GPU.

Generation-request execution workflow

With the model has been loaded, initialized, and set up in vLLM, it’s ready to serve. Let’s look at a high-level workflow for executing generation requests (pictured in four steps in Figure 8-4).

Figure 8-4. The vLLM generation-request execution workflow

When you run llm.generate(prompts, sampling_params, internally, here is what vLLM does:

The Processor component validates and preprocess the raw inputs into a Request object, including tokenizing the input prompts.

LLMEngine runs an execution loop, which calls the EngineCore repeatedly to process requests. The EngineCore lets the Scheduler decide the next batch of requests to run and the tokens to be processed. Scheduler also applies many optimizations in this step, such paged attention and continuous batching.

EngineCore passes the scheduled token (SchedulerOutput) from Scheduler to MultiProcessExecutor.

MultiProcessExecutor passes the request to the worker process and delegates the GPU worker to run the model forward pass to calculate the given tokens. This is where the actual model execution happens.

EngineCore passes the model output to the output processor, which creates the final response output to the customer.

From this serving workflow, you can see that the actual model execution and model-specific optimization are handled by the GPUWorker, while the general LLM-serving planning and optimization are orchestrated by the Scheduler. Next, let’s take a deeper look at how the Scheduler works.

Scheduler deep dive

In general, the Scheduler functions as the central traffic controller for executing generation requests in vLLM’s inference pipeline. We will first examine the key considerations that shape the vLLM Scheduler, and then dive into its request scheduling workflow:

Request Resource Orchestration

The Scheduler orchestrates the entire lifecycle of a request, from its arrival to completion. It receives incoming requests, organizes them into separate WAITING and RUNNING queues, and makes dynamic decisions about which requests to dispatch based on available computational resources such as GPU memory, KV cache blocks, and token budgets.

The Scheduler also tackles complex resource-allocation challenges, including speculative decoding, prefix caching, and multimodal inputs. Its goal is to maximize throughput while maintaining fairness across concurrent requests, ensuring an efficient and balanced use of system resources.

Token-Level Resource Allocation and Scheduling

One of vLLM’s key strengths is that it uses a unified token-based scheduling approach, rather than separating the Prefill and Decode phases. Unlike a request-based scheduler, which treats each request as a whole unit, vLLM’s Scheduler operates on a token-by-token basis. At each scheduling step, it decides how many tokens each request can process, while respecting global constraints such as maximum batch size and GPU memory limits.

Compared to a request-based scheduling strategy, token-based scheduling offers finer-grained control. It balances competing demands across requests by allocating token budgets, managing KV cache blocks for attention states, and supporting encoder inputs in multimodal models. This approach ensures the total computational load stays within system limits while still maximizing opportunities for parallel execution.

Optimization Integration Hub

The Scheduler acts as the central integration point for model-agnostic performance optimizations. It coordinates techniques such as prefix caching (reusing previously computed states), speculative decoding (predicting future tokens), chunked prefill (efficiently processing long sequences), and distributed KV cache transfers (enabling multi-GPU execution).

Crucially, the Scheduler makes adaptive decisions on when and how to apply these optimizations. Its choices are guided by request characteristics, resource availability, and overall system state, ensuring that each optimization improves performance without introducing conflicts or inefficiencies.

Dynamic Load Balancing

The Scheduler monitors system resources and requests characteristics to make real-time decisions that balance latency and throughput. It adapts to dynamic events—arrivals, completions, preemptions, and resource limits—by reevaluating which requests to run, how many tokens to process, and when to apply optimizations. Through efficient queue management and resource allocation, it coordinates with workers to sustain optimal performance.

Request Lifecycle Management

The Scheduler monitors each request throughout its lifecycle—from arrival in the WAITING queue through execution in the RUNNING queue to final completion.

Scheduler also applies advanced scheduling policies, such as first-come-first-served (FCFS) or priority-based ordering, to determine execution order. When resources are limited, the Scheduler can preempt lower-priority requests, reallocating capacity to higher-priority tasks. It also manages smooth transitions between request states, ensuring both efficiency and fairness in execution.

Request Scheduling Workflow

The core task of the Scheduler is to determine the appropriate number of tokens for each incoming generation request and pass them to the model for a forward pass. To fully utilize model and hardware capacity while maintaining a balanced user experience, the Scheduler applies a variety of model-agnostic optimizations during this token selection process. The overall request-scheduling logic is illustrated in Figure 8-5.

Figure 8-5. vLLM internal requests scheduling logic

From Figure 8-5, you can see the vLLM scheduler begins its scheduling cycle by building an internal Schedule State, which gathers all relevant requests: newly arrived requests, previously paused requests that need to be resumed, currently running requests, and any requests that were preempted earlier. During this initialization phase, the scheduler also updates its accounting of available resources, such as how many decoding tokens remain within the current token budget and how much encoder capacity is still available for multimodal workloads. This initial state frames all decisions that follow.

Once initialization is complete, the scheduler begins assigning compute in two prioritized waves. Running requests are handled first, because they already occupy KV cache blocks and have active generation timelines. For each running request, the scheduler determines how many new tokens must be generated, verifies encoder constraints when multimodal inputs are involved, and checks whether there is enough available KV cache to continue execution.

At this stage, vLLM applies several optimizations that significantly reduce latency and improve throughput. Chunked prefill allows partial processing of long prompts without requiring the entire prompt to be encoded in a single step. Prefix caching enables the scheduler to reuse KV cache blocks that were computed earlier for identical prefixes across different requests. Meanwhile, speculative decoding generates draft tokens ahead of time, enabling faster generation if those predictions later match the model’s verified output. If at any point resources become constrained, the scheduler may preempt lower-priority work to ensure fairness and consistency.

After all running requests have been handled, the scheduler turns to the waiting queue. Any requests that can fit within the remaining token and encoder budgets are activated and incorporated into the current execution batch. Although waiting requests receive lower priority, they still benefit from the same optimizations applied to running requests whenever applicable.

When request assignment is complete, the scheduler enters a post-processing phase that performs additional coordination work necessary for the upcoming decoding step. In this phase, the scheduler tracks which LoRA adapters must be active for each request, prepares encoder inputs for multimodal tasks, and finalizes any draft-token predictions that will be used during speculative decoding. The goal of this stage is to consolidate all scheduling decisions into a clean, executable plan.

Finally, the scheduler assembles a SchedulerOutput object that summarizes the results of the entire cycle. This output lists the requests that are newly scheduled, the number of tokens allocated to each request, and the total number of scheduled tokens across the batch. It also includes metadata needed by the model executor, such as KV cache assignments, prepared encoder inputs, or multimodal routing information. The model executor then consumes this output to run the actual forward passes that produce new tokens.

In addition, to support a model-agnostic request–token scheduling process, the vLLM Scheduler operates at the token level. It iterates through each request—first from the running queue, then from the waiting queue—and determines how many tokens from which request to include in the next forward pass, while ensuring the total token count remains within the model’s limits.

At the implementation level, the Scheduler seeks to minimize the gap between num_computed_tokens—the number of tokens already processed—and num_tokens_with_spec—the total tokens to be processed, including prompt, output, and speculative tokens. This goal of closing the gap drives the chain of optimization strategies applied during scheduling. Here’s how that looks in the code:

while req_index < len(self.running) and token_budget > 0:
    request = self.running[req_index]
    num_new_tokens = (request.num_tokens_with_spec +
                      request.num_output_placeholders -
                      request.num_computed_tokens)

Separate prioritization and optimization

In vLLM scheduling, the request queues determine the order of request processing, while the comparison between num_computed_tokens and num_tokens_with_spec dictates how many tokens each request is allowed to execute.

This clear separation of request prioritization and token-level scheduling allows the Scheduler to combine diverse request-prioritization strategies with LLM execution optimizations in a unified framework, maintaining fairness across concurrent requests while ensuring efficient use of system resources.

Applying Model Optimization Techniques

During scheduling, vLLM integrates several of the optimization techniques we introduced in Chapters 6 and 7, including chunked prefill, prefix caching, guided decoding (via grammar-constrained finite-state machines), disaggregated prefill/decoding (P/D), and more. To illustrate how these work in practice, let’s walk through two quick examples:

Chunked prefill

The Scheduler ensures that the number of new tokens executed for a request does not exceed the configured prefill chunk size–scheduler_config.long_prefill_token_threshold:

if (0 < self.scheduler_config.long_prefill_token_threshold \
                  < num_new_tokens):
   num_new_tokens = self.scheduler_config.long_prefill_token_threshold

Prefix caching

The Scheduler handles prefix caching to reuse computed KV cache blocks from local or remote cache:

# Get already-cached tokens.
if request.num_computed_tokens == 0:
    # Get locally-cached tokens.
    new_computed_blocks, num_new_local_computed_tokens = \
          self.kv_cache_manager.get_computed_blocks(
                           request)

    # Get externally-cached tokens.
    if self.connector is not None:
        num_external_computed_tokens, load_kv_async = (
            self.connector.get_num_new_matched_tokens(
                request, num_new_local_computed_tokens))

We encourage you to explore other optimization techniques and their implementations—such as guided decoding and prefill/decoding (P/D)— further on your own. A good starting point is the Scheduler.py code implementation. Additionally, Aleksa Gordić’s blog post “Inside vLLM: Anatomy of a High-Throughput LLM Inference System” is an excellent resource.

vLLM’s Layered Optimization Strategy

One of the most important optimization philosophies in vLLM is that optimization must happen at the right level.

LLMs evolve rapidly—new architectures, novel attention mechanisms, and hardware-driven execution tricks appear almost monthly. A system like vLLM cannot hard-code its optimizations for one model family or hardware type; instead, it must provide a flexible yet unified framework that supports arbitrary architectures and model sizes, while still delivering system-wide throughput, fairness, and resource efficiency.

Goals and Challenges

Since vLLM’s overarching goal is to maximize throughput while preserving fairness across concurrent requests, it needs to handle many users, diverse workloads, and variable model sizes without starving smaller requests or overloading GPU memory. The challenge is that no single layer of the system can address all optimization needs. Some optimizations apply universally across models (such as, batching and caching), while others are tied to model architecture (like, attention kernels) or even specific hardware instructions (for example, tensor cores).

To resolve this tension, vLLM employs a layered optimization strategy, organizing optimization responsibilities across four main levels:

Scheduler: System-wide, model-agnostic optimizations

The Scheduler is responsible for fairness, efficiency, and scalability at the system level.

Model Executor: Model architecture-specific optimizations

While the Scheduler remains model-agnostic, the Model Executor understands the details of each model architecture. For example, it applies fused attention kernels for transformer-based models or specialized operators for multimodal encoders. This level isolates architecture-aware optimizations so they can evolve independently from system-level scheduling.

Model Layer: Component-specific optimizations

At the model architecture’s layer level (for example, the attention layer and feed-forward block), optimizations are tailored to computational bottlenecks. Techniques such as KV cache reuse, flash attention, and layer-wise operator fusion happen here. This design ensures that optimizations that target a specific subcomponent don’t leak into the system-wide scheduling logic.

CustomOps: Hardware-specific optimizations

Finally, CustomOps capture optimizations for the underlying hardware, such as CUDA kernels, tensor-core acceleration, or quantized operators. By isolating them here, vLLM can take advantage of new GPU capabilities and accelerators without changing its higher-level scheduling or model logic.

This layered strategy avoids cluttering the Scheduler with model-specific details while still allowing deep specialization where it matters. It also makes vLLM future-proof: as new architectures or hardware arise, optimizations can be slotted into the right layer without redesigning the entire system.

TensorRT-LLM

TensorRT-LLM is NVIDIA’s open-source library for high-performance LLM inference on NVIDIA GPUs. From a model checkpoint, it builds highly tuned TensorRT (TRT) engines and provides Python/C++ runtimes with modern serving capabilities, including: in-flight batching (TRT-LLM version of continuous batching), paged KV cache, speculative decoding, multi-precision quantization (FP8/FP4/INT4/INT8), and tensor/pipeline parallelism. It integrates tightly with the NVIDIA serving ecosystem—most notably Dynamo and Triton.

Let’s look at a quick example with TensorRT-LLM’s high-level LLM API:

llm = LLM(model="Qwen/Qwen3-7B")

# Sample prompts.
prompts = [
   "Hello, my name is",
   "The capital of France is",
   "The future of AI is",
]

# Create a sampling params.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

# Run model generation request
for output in llm.generate(prompts, sampling_params):
   print(
       f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
   )

Note

A core objective of TensorRT-LLM is to demonstrate the maximum practical performance of NVIDIA GPUs for LLMs and to illuminate how to fully exploit Tensor Cores and CUDA kernels.

TensorRT-LLM is best for organizations that have standardized on NVIDIA’s hardware and serving tech stack and are seeking best-in-class efficiency and throughput in production.

SGLang

SGLang is a newer entrant, targeting structured-generation and agent applications. SGLang is an open-source, high-performance serving framework for LLMs and vision-language models (VLMs). It co-designs a fast backend runtime (kernels, caching, scheduling) with a flexible frontend language and API (which is OpenAI-compatible and native) to make generation faster and more controllable.

SGLang’s headline features (available on GitHub) include RadixAttention (for prefix/KV reuse across calls), continuous batching, paged attention/KV, speculative decoding (EAGLE-2/3), chunked prefill, structured outputs (JSON/regex/EBNF), multi-LoRA batching, and parallelism modes (tensor/pipeline/expert/data). A quick example:

# load Qwen3 model
llm = sgl.Engine(model_path="Qwen/Qwen3-7B")
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

sampling_params = {"temperature": 0.8, "top_p": 0.95}

# Run model generation requests
outputs = llm.generate(prompts, sampling_params)
for prompt, output in zip(prompts, outputs):
    print("===============================")
    print(f"Prompt: {prompt}\nGenerated text: {output['text']}")

SGLang targets a broader range of hardware than TensorRT-LLM does. Its official documentation and examples cover NVIDIA (H100/Blackwell), AMD Instinct, CPU, TPU, Jetson Orin, and Ascend deployments, making it attractive for multi-vendor fleets.

Functionally, SGLang can be viewed as a peer to vLLM: it emphasizes infrastructure portability, broad open-model support, aggressive performance optimizations (for example, continuous batching, speculative decoding, and prefix/KV reuse), scale-out routing, and strong support for agentic, multi-step prompts.

In practice, SGLang’s performance is competitive—sometimes better on certain workloads—while vLLM currently enjoys a larger open-source community and ecosystem.

Llama.cpp

Llama.cpp is a lean, open-source C/C++ inference stack built to run modern open-weight LLMs efficiently on almost any machine—from laptops and workstations to on-prem servers and edge devices. Models are packaged in the GGUF format (typically converted from Hugging Face checkpoints with the repo’s scripts). You can serve them via a built-in, OpenAI-compatible HTTP server or drive them from small CLI tools for experiments and benchmarking.

The project emphasizes minimal dependencies, fast startup, aggressive integer quantization (for example, 8/6/5/4-bit), and a portable set of CPU/GPU backends (CPU with SIMD, plus optional accelerators such as Metal, CUDA/ROCm, or Vulkan).

What distinguishes llama.cpp is its “runs anywhere” ethos and tiny operational footprint. Its primary goal is to deliver state-of-the-art local inference with minimal setup and solid performance across a wide range of hardware, including low-cost or low-profile devices, both offline and in the cloud. By contrast, the other three serving frameworks covered here—vLLM, TensorRT-LLM, and SGLang—optimize first for high-throughput, low-latency serving on datacenter GPUs, such as multi-tenant batching, advanced scheduling, and large-scale deployment features.

Llama.cpp prioritizes portability, simplicity, and cost efficiency over absolute peak throughput, making it ideal for local development, privacy-sensitive workloads, edge deployment, and budget-constrained environments. However, it still offers optional GPU acceleration when available.

Let’s see a quick example for running the Qwen3 model on a CPU device with llama.cpp:

from llama_cpp import Llama
# Run Qwen3 model on CPU with LlamaCPP

# load Qwen model (in GGUF format) from HuggingFace
llm = Llama.from_pretrained(
   repo_id="Qwen/Qwen3-8B-GGUF",
   filename="*Q8_0.gguf",
   verbose=False
)

# Run LLM generation with highlevel API
output = llm(
     "Q: Name the planets in the solar system? A: ", # Prompt
     max_tokens=32, # Generate up to 32 tokens
     stop=["Q:", "\n"], # Stop generating just
          # before the model would generate a new question
     echo=True # Echo the prompt back in the output
)

# Example for chat completion API
output = llm.create_chat_completion(
     messages = [
         {"role": "system", "content": "You are an assistant who perfectly describes images."},
         {
             "role": "user",
             "content": "Describe this image in detail please."
         }
     ]
)

Because llama.cpp brings LLM serving to low-profile devices (CPUs and edge hardware), it sees strong adoption in three following scenarios:

Local development, where you can run the built-in OpenAI-compatible server on your machine and reuse existing clients (for example, RAG and search systems)

Private and on-prem assistants, keeping all data inside your VPC or device boundary

Edge inference on laptops, desktops, Apple Silicon, and small servers—enabled by aggressive 2 to 8-bit quantization and portable CPU/GPU backends

If you want to expose llama.cpp through a REST API call, you can use Ollama, a higher-level framework that wraps around llama.cpp (and sometimes other backends, like Mistral.cpp or RWKV runners) to make local LLM use simple, consistent, and developer-friendly.

Note

Not every LLM application needs high-throughput serving. When you run inference locally (on-device or on-prem edge), the goals shift:

You want to optimize for latency and responsiveness, not fleet-wide tokens per second.

Low concurrency (often a single user) makes large batching and complex schedulers less valuable.

Footprint and cost (in memory, compute, and power) become primary constraints.

Privacy and offline reliability are first-class requirements.

Llama.cpp fits this profile: it can run open-weight models directly on CPUs, Apple Silicon, and small GPUs, and exposes an OpenAI-compatible server for drop-in integration. The result is a super-low-cost, low-ops serving option for local development, private and on-prem assistants, and edge deployments—without per-token cloud charges or data leaving the device.

Selecting the Right Framework

The general idea is to choose the framework that best matches your Service Level Objectives (SLOs), hardware, and operational reality—not the one with the flashiest benchmark. In other words: workload first, not hype.

The evaluation approach we generally recommend is:

Start with SLOs, not features. Write down your targets for latency (TTFT, p95/p99), throughput (tokens/sec/QPS), cost per token, quality constraints (structured JSON, safety), and availability.

Analyze the real prompts from your use case, make it clear whether it is prefill-heavy or decode-heavy; what’s the context length; does it include tool calls or multi-turn chains.

Make it an apples-to-apples comparison using the same model, dtype/quantization, max sequence, batch/concurrency, and streaming settings across contenders.

Measure operability, using metrics such as cold-start time, observability, autoscaling behavior, multi-tenancy fairness, upgrade friction, and failure modes.

Account for hardware and vendor lock-in. If you use multiple vendors (for instance, NVIDIA/AMD/CPU/TPU/edge), portability weighs heavily.

Plan for change. Model swaps and new decoding tricks happen weekly, so opt for frameworks you can update without major surgery.

Following our own advice, we’ll conclude our library recommendation as follows:

If you want the fastest path to production, broad model coverage, and a Python-native workflow, choose vLLM. It offers strong baseline performance, continuous batching, and a large community and ecosystem.

If your workload involves agentic or multi-step workflows, strict JSON/regex (grammar-constrained) outputs, and multi-vendor portability, pick SGLang. You get prefix/KV reuse, speculative decoding, and a router for scale-out clusters.

If you’re all-in on NVIDIA and its serving stack and want peak tokens per dollar at scale, start with TensorRT-LLM. First-class Triton and Dynamo integration and deep CUDA/Tensor Core optimizations make it the efficiency leader on NVIDIA hardware.

If you need local, on-prem, or edge serving with a tiny footprint, low cost, and privacy by default, use llama.cpp. GGUF quantization (2 to 8-bit) and portable CPU/GPU backends make it ideal for laptops, desktops, and small servers.

Keep an open mind when evaluating serving frameworks

LLM serving evolves fast—new architectures, kernels, schedulers, quantization schemes, and runtime features land every month. Frameworks must adapt just as quickly to incorporate these innovations.

As serving engineers, we evaluate and reassess different serving frameworks every 3 to 6 months, and select the option that best fits our current SLOs and hardware reality. We also need to build a framework abstraction layer and have a clear exit plan that lets us swap frameworks without rewriting our apps.

The goal isn’t to crown a permanent winner—it’s to remain flexible so our system can ride the next wave, not be capsized by it.

Summary

In this chapter, we outlined why general-purpose ML serving stacks fall short for LLMs—chiefly because they lack LLM-specific feature support such as token-level scheduling, KV-cache management, long-context memory handling, and streaming-first execution.

We then took a deep dive into vLLM, covering its architecture, request scheduling and prioritization, and layered optimization strategy, which separates system-wide scheduling from model- and kernel-level specialization. We rounded out the survey with three complementary frameworks: TensorRT-LLM for pushing peak efficiency on NVIDIA GPUs, llama.cpp for lightweight local/on-prem/edge serving, and SGLang for agentic, multi-step workflows and structured outputs across diverse hardware.

The core takeaway we offer you is that your framework choice is contextual. Many teams begin with vLLM for online serving and llama.cpp for local development, but specific requirements—like edge deployment, NVIDIA-centric optimization, strict JSON/grammar outputs, or agent pipelines—may point you to TensorRT-LLM or SGLang instead. As the ecosystem evolves, it’s a good idea to stay framework-agnostic. In practice, this means you should do periodic framework research, prefer portable interfaces, and be ready to swap components so your serving stack can adapt to changing business and technical needs.

About the Authors

Chi Wang has over 17 years of experience in the tech industry, with a particular focus on artificial intelligence and distributed systems. For the past 8 years, Chi has been a key contributor at Salesforce’s Einstein AI group, where he leads the development of AI platforms and infrastructures that support millions of Salesforce customers and power hundreds of AI features. Currently, as the Director of Engineering, Chi oversees two critical teams: one focused on model serving and optimization solutions, and the other on data science environments. Chi also filed 12 patents in areas such as dataset management, model serving and optimization, data access authorization, and networking management. In addition, he holds an Artificial Intelligence Graduate Certificate from Stanford University, which he completed in 2020.

Peiheng Hu is an accomplished machine learning engineer with over 10 years of industry experience and expertise in building robust large-scale AI driven systems on the cloud. He holds a Master of Science in Computational Science and Engineering from Harvard University and a Bachelor of Science in Industrial Engineering Operations Research from Georgia Institute of Technology. Peiheng currently serves as a Principal Member of Technical Staff and ML Engineer at Salesforce, where he leads teams in developing cutting-edge machine learning inferencing solutions, including launching Salesforce’s only unified ML inferencing solution which now handles thousands of requests per second and a novel automated model optimization framework for Large Language Models (LLMs). His work has significantly enhanced model inference performance, scalability and cost-efficiency, saving millions in hardware expenses.