Hello, this is the Product Development Department at Qualiteg Inc.

Today we would like to introduce PagedAttention, a technology that is indispensable for commercial LLM serving.

Introduction

PagedAttention is an extremely important technology for our company.

By taking advantage of PagedAttention, you can improve GPU memory utilization during text generation with LLMs, and correspondingly increase the number of requests that can be handled concurrently per GPU.

Because we develop and offer a commercial LLM serving platform called ChatStream, LLM serving (that is, providing text generation) under heavy load from many concurrent users is a challenge that sits squarely at the center of what we do.

Before PagedAttention came along, conventional parallel generation was a constant battle against the KV cache, the large per-request GPU memory consumption.

(When you call a Transformers model directly, the KV cache appears as past_key_values.)

In other words, it is memory consumed at inference time on top of the model parameters themselves.

This determines the upper limit on the number of requests that can be handled concurrently. That limit is the very edge at which inference-time memory consumption does not exceed the GPU's onboard memory, and before reaching it, generation requests must be offloaded to another GPU node. Raising this request limit translates directly into lower inference cost per request (or per token).

We had already achieved stable parallel generation under multi-user access through our own techniques, but PagedAttention, which we introduce here, is an even more efficient method. At the same time, vLLM, the serving engine that implements it, offers not only PagedAttention but also model parallelism across compute nodes using tensor parallelism, making it a highly promising base engine for commercial inference environments.

Our ChatStream supports both algorithms: conventional (*) parallel generation using Transformers, and the more efficient parallel generation using PagedAttention, so that we can offer a high degree of customizability and meet a wide range of needs.

* We continue to support the conventional approach because it lets us leverage our assets of numerous preset sampling algorithms and original sampling algorithms, and because for the latest GPUs with large memory, or for local LLM serving that is not under heavy load, the conventional approach still delivers sufficient performance.

The Impact of PagedAttention

Around the summer of 2023, the paper we introduce here, "Efficient Memory Management for Large Language Model Serving with PagedAttention," and a library called vLLM were released, marking an important breakthrough in LLM serving.

The paper is available here:
https://arxiv.org/abs/2309.06180

In Japan, it did not seem to attract much attention when it was first published (perhaps because few companies make LLM serving their core business, as we do), but for us it was a major shock in many respects. Our engineering team had just rolled up its sleeves, thinking that the overly straightforward KV cache implementation in Transformers was easy to understand but could surely be done a little more cleverly, when suddenly it felt as if we had been hit by a 100-ton hammer. Today, we are developing and trialing even more efficient LLM serving methods built on top of PagedAttention.

Main Article

So in this post, we would like to explain this paper, which had such a strong impact on our own products, in language that is as plain as possible, while weaving in our own perspective on the challenges involved.

Since we prioritize a high-level understanding, some parts are deliberately simplified. If you want to understand the details, we recommend consulting the original paper or the vLLM repository (https://github.com/vllm-project/vllm).

(Personally, my understanding of the parts that the paper only touches on briefly deepened considerably by reading through the source code.)

In a Nutshell, What Problem Does PagedAttention Solve?

The conventional problems

• Problem 1: There is a technique called the KV cache that speeds up LLM computation by caching the outputs of past layers. Reusing past computation results speeds up LLM computation, but in the conventional approach, the GPU memory (tensor) used for this KV cache was assigned "in advance" as a "contiguous region" sized for the "maximum number of tokens" the model can handle.
  In other words, with the conventional method, when generating text for a given request, extra memory was reserved even when there was no real need to consume that much. As a result, the KV cache used a large amount of memory, and a great deal of memory was wasted on every inference.
• Problem 2: In addition, in LLM inference scenarios, batch computation, which is needed to fully benefit from massively parallel computation, the GPU's inherent strength, could not be exploited effectively. (*)

* We explain why this problem existed later in the article.

How PagedAttention solves and improves on these problems

Here is a concise explanation of how PagedAttention solved the problems above.

[Regarding Problem 1]

• 1-1 Where the conventional approach reserved KV cache memory for the "maximum number of tokens," PagedAttention reserves only "as much as needed."
• 1-2 Conventionally, the tensors used as the KV cache had to be stored in a "contiguous region" of GPU memory, but PagedAttention divides the KV cache into "blocks," eliminating the need to place it in contiguous memory space.
• 1-3 Conventionally, KV cache memory was reserved "in advance," so depending on the request to the LLM, there was surplus memory that ended up unused. PagedAttention can add or remove blocks "as needed," so there is no longer any need to reserve memory in advance.

[Regarding Problem 2]

• 2-1 In conventional batch processing, when multiple texts are input, the lengths of the token sequences serving as prompts have to be aligned, which produced unnecessary padding. PagedAttention introduces a special GPU kernel that eliminates the need for padding at input and output.
• 2-2 Conventionally, unless requests arrived at exactly the same time, an incoming request had to wait until the previous request finished. With PagedAttention, this "per-request" synchronization becomes "per-iteration," allowing requests to be swapped in and out at a fine granularity, which shortens waiting times.

That was not quite a nutshell, but roughly speaking, that is the picture.

Characteristics of Autoregressive Models and the Difficulty of Estimating Required Memory

An LLM is an autoregressive model, which generates a new token by feeding back in the sequence of tokens generated so far.

Each of these repetitions is called an iteration, and the target text is generated through iterations. In this repeated computation process, computation time can be shortened by reusing results computed in the past. This is where the KV cache comes in. Put simply, a KV cache is prepared for as many tokens as there are.

The problem here, and it is a characteristic of autoregressive models, is that from the start of text generation to the end, there is no way to know in advance how many iterations will be needed.

Because it cannot be known in advance, the question arises of how much memory should be reserved in advance for the KV cache.

The simplest strategy is to reserve "the maximum token size the model can handle" at the start of generation. We named this the "naive implementation," and that strategy was in fact what was implemented. Being too naive, it consumed a large amount of KV cache memory during generation.

That is fine at the level of laboratory experiments, but for those of us hosting commercial LLMs, it becomes a major problem.

Deepening Our Understanding of PagedAttention

Now that we have refreshed our memory on autoregressive models, let's look at concrete examples to deepen our understanding.

1-1 Where the conventional approach reserved KV cache memory for the "maximum number of tokens," PagedAttention reserves only "as much as needed."

Let's start with the point above.

Here we will use a concrete example: text generation with rinna/japanese-gpt2-small, a Japanese model based on GPT-2.

First, let's look at the context size of rinna/japanese-gpt2-small.

We check it with the following code:

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt2-small", use_fast=False)

model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt2-small")
max_position_embeddings = model.config.max_position_embeddings

print("Model context size (maximum sequence length):", max_position_embeddings)

The result was 1024. This means the model can generate up to 1024 tokens.

We will generate text like the following with this model.

Consider the example where, generating the continuation of 「夏目漱石の代表作は」 ("Natsume Soseki's most famous work is"), the model outputs 「吾輩は猫である」 ("I Am a Cat").

(In reality, since this is GPT-2, it will not produce such a convenient output.)

In this case, the "naive implementation" behaves as follows.

Memory was reserved for the full context size of 1024, yet only 13 tokens were generated, so the remaining 1011 tokens' worth went unused, in other words, wasted.

Checking this in an implementation looks like the following.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt2-small", use_fast=False)
tokenizer.do_lower_case = True

model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt2-small")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

prompt = "夏目漱石の代表作は"
inputs = tokenizer(prompt, return_tensors='pt', padding=True).to(device)
input_ids = inputs['input_ids']

max_length = 20
outputs = input_ids.clone()  # Clone the input IDs and keep them for the output

use_cache = True  # If set to False, the past cache cannot be used and the computation goes wrong

# First output generation and retrieval of the KV cache (past_key_values)
out = model(input_ids=input_ids, use_cache=use_cache) 
logits = out.logits
past_key_values = out.past_key_values  # Get the KV cache from the first output
last_token_logits = logits[0, -1, :]
token_id = torch.argmax(last_token_logits).unsqueeze(0).unsqueeze(0)  # Convert the scalar to a [1,1] tensor
outputs = torch.cat((outputs, token_id), dim=1)

# Second and subsequent output generation. Generate using the KV cache (past_key_values) and update the KV cache
for idx_itor in range(1, max_length):
    out = model(input_ids=token_id, past_key_values=past_key_values, use_cache=use_cache)
    logits = out.logits
    past_key_values = out.past_key_values  # Explicitly update past_key_values. Some models update it internally without an explicit update.
    last_token_logits = logits[0, -1, :]
    token_id = torch.argmax(last_token_logits).unsqueeze(0).unsqueeze(0)  # Convert the scalar to a [1,1] tensor
    outputs = torch.cat((outputs, token_id), dim=1)  # Append the new token ID to outputs

# Display the generated text
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"{generated_text}")

Setting use_cache = True below makes the model keep the KV cache.

out = model(input_ids=input_ids, use_cache=use_cache) 

Furthermore,

past_key_values = out.past_key_values

lets you retrieve the KV cache.

out = model(input_ids=token_id, past_key_values=past_key_values, 

The line above is the second and subsequent generation steps. By passing past_key_values as an argument, that is, the results computed so far, the corresponding computation cost is skipped, and only the cost attributable to the newly input token is incurred. That is why only the new token token_id is fed in here.

Incidentally, sampling is done simply by taking the argmax.

token_id = torch.argmax(last_token_logits)

Now, this runs on Hugging Face Transformers, but does it really expand and reserve memory up to the context size at the time of the first input, as the paper claims? The answer is no.

In practice, memory usage grows as the sequence gets longer.

1-2 Conventionally, the tensors used as the KV cache had to be stored in a "contiguous region" of GPU memory, but PagedAttention divides the KV cache into "blocks," eliminating the need to place it in contiguous memory space.

Regarding the point above:
In the conventional implementation, the KV cache tensor data is allocated "in the ordinary way," so PyTorch tensors are allocated and used at the size of the KV cache as is. Inevitably, the tensor data is then laid out sequentially in GPU memory space, taking the form shown in "Before" below. The proposed method, on the other hand, divides the KV cache into blocks in advance, as shown in "After" below, and thereby escapes the constraint of having to place it contiguously in GPU memory.

(However, if the block size is too large, the same problem, fragmentation, occurs after all, so a block size of around 16 to 256 is considered best.)

What about the Hugging Face Transformers implementation on this point?

On this point, the KV cache in Transformers is still an ordinary PyTorch tensor, so the answer is yes, in the sense that a contiguous region of GPU memory does get tied up.
(The implementations of Transformers and other components that serve as inference engines evolve day by day, and obvious problems are not left unaddressed. To understand how far a library has evolved, it is best to keep up with the latest version at the code level.)

1-3 Conventionally, KV cache memory was reserved "in advance," so depending on the request to the LLM, there was surplus memory that ended up unused. PagedAttention can add or remove blocks "as needed," so there is no longer any need to reserve memory in advance.

The point above is already covered by the explanation so far, so we will omit further explanation.

There are logical blocks and physical blocks, and the two are linked by a block table

One more technical point worth explaining: in the block design, blocks are not placed directly in GPU memory. Instead, they are first abstracted as logical blocks, sequences are reconstructed and blocked at the logical block level, and the actual storage destination is GPU memory at the physical block level. Logical and physical blocks are linked by a block table. You can think of it as something like a cross-reference table in a database.

Next, let's look at Problem 2.

[Regarding Problem 2]

2-1 In conventional batch processing, when multiple texts are input, the lengths of the token sequences serving as prompts have to be aligned, which produced unnecessary padding. PagedAttention introduces a special GPU kernel that eliminates the need for padding at input and output.

2-2 Conventionally, unless requests arrived at exactly the same time, an incoming request had to wait until the previous request finished. With PagedAttention, this "per-request" synchronization becomes "per-iteration," allowing requests to be swapped in and out at a fine granularity, which shortens waiting times.

For these, let's start with 2-1.

LLM Serving and Batch Processing

First, let's consider why GPU batch computation does not work well in LLM serving.

To think about that, let's first review batch processing when the GPU performs the LLM's forward computation (that is, inference).

Batch processing means bundling input token lists into a batch, computing on them, and outputting a bundle of output token lists.

To deepen our understanding, let's again think through a real example.

Batch processing problem 1: The padding problem

As before, consider the scenario where we input 「夏目漱石の代表的な作品は」 ("Natsume Soseki's representative work is") to the LLM
and the output is 「吾輩は猫である」 ("I Am a Cat").

With a single, non-batched input, the picture looks like this.
(Since this is a schematic diagram, special characters and end-of-sequence tokens are omitted.)

After 「夏目漱石の代表作は？」 ("What is Natsume Soseki's most famous work?") is tokenized and fed into the LLM, 「吾輩は猫である」 ("I Am a Cat") is returned as the output (including the input prompt).

The numbers 1 through 6 are generation iterations: 6 tokens are generated in 6 iterations, yielding 「吾輩は猫である」.

A schematic view of the generation iterations at batch size 1 looks like the figure below.

Written with Transformers, this looks like the following.

import torch

# Load the model and tokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt2-small", use_fast=False)
tokenizer.do_lower_case = True

model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt2-small")

if torch.cuda.is_available():
    model = model.to("cuda")

model.eval()

prompts = ["夏目漱石の代表作は"]

inputs = tokenizer(prompts, return_tensors='pt', padding=True).to(model.device)

outputs = inputs.input_ids

max_length = 20

for _ in range(max_length):
    with torch.no_grad():
        outputs_logit = model(input_ids=outputs)

    next_token_logits = outputs_logit.logits[:, -1, :]  # From the model output, get the logits at the last token position
    next_tokens = torch.argmax(next_token_logits, dim=-1, keepdim=True)  # Simple sampling: take the most probable token
    outputs = torch.cat([outputs, next_tokens], dim=-1)  # Append the token to the output

# Display the generated text
for i, output in enumerate(outputs):
    print(f"{i + 1}: {tokenizer.decode(output, skip_special_tokens=True)}")

What happens with batch input?

Batch input to an LLM means that multiple texts can be generated in parallel in a single request.

Consider performing the following two generations at the same time:

1. 「夏目漱石の代表作は」 ("Natsume Soseki's most famous work is") → 「吾輩は猫である」 ("I Am a Cat")
2. 「君達は」 ("How do you") → 「どう生きるか」 ("live?", i.e., "How Do You Live?")

Drawn as a diagram, this looks like the following.

You can see that PAD appears here.

As the figure shows, the sizes (input token sizes) of the two prompts must be aligned. The process of filling in the missing tokens on the 「君達は」 side is called padding. The observation that this padding is wasteful was the concern behind problem 2-1.

(Although omitted in the code above, it is possible to deliberately ignore the padded portion by explicitly specifying an attention mask.)

Batch processing problem 2: The timing problem

The other observation is this: when you consider LLM serving in practice, do text generation requests ever conveniently arrive at exactly the same time?

The ideal case (computationally efficient for the GPU and fast responses for every user) is when requests happen to arrive simultaneously and can be generated as a batch.

The following shows the case where four users send requests at exactly the same time. The GPU performs generation at batch size 4. Time flows to the right.

Because this is batch inference, it is very convenient both for the GPU and for the UX, but such a convenient situation almost never occurs except under very heavy concurrent multi-user access.

In other words, moderate concurrent access is usually the norm, and request processing timings never coincide perfectly.

In that case, with a "naive" implementation, if request timings are staggered, other requests have to wait while one request is being processed.

The following shows Request 1 being generated as a single request (batch size 1). User 1 gets the generated output, but User 2 through User 4 are left waiting until their own requests are processed. (Of course, this is because we are simplifying the discussion to a single GPU.)

The Conventional Approach in Qualiteg's ChatStream

Here we would like to briefly introduce our own approach to this problem.
To address this problem, our ChatStream focused on the following points to make concurrent access handling more efficient.

1. Split requests into iteration units

2. Single batch, but allow other requests to interrupt at iteration granularity

With this, regardless of request length, we handled concurrent access by monitoring and controlling the number of requests currently being processed and the memory usage.

We call this method Qualiteg classic hf-transformer, or QCHT.

This improved concurrency and responsiveness, but there was still a problem: once again, the KV cache. With this method too, the longer the sequence generated by a request, the more memory it uses.

Naturally, if too much memory is used, processing crashes, so memory consumption has to be kept within a reasonably safe level.

Here, if memory is estimated conservatively, one possible method is to take the model's context size, or the maximum context size decided by the application, as the upper bound, and estimate for each request the amount of memory that may be used even if the maximum context size (sequence length) is reached.
In this case,

total GPU memory = memory occupied by the model data itself + number of concurrent users × maximum context size × KV cache memory per token + other memory.

Setting aside the other memory for a moment, it can be calculated as

total GPU memory = memory occupied by the model data itself + number of concurrent users × maximum context size × KV cache memory per token

.

We plan to explain how to compute actual values in a separate article, but with this method, the problem is that you cannot achieve a high number of concurrent users unless you put a model with a small parameter size on a GPU with large memory.

So we used the following methods in a hybrid fashion:

1. Decide on a minimum amount of reserved memory
2. Monitor memory usage during generation and stop accepting requests when it reaches a specified amount (*)
3. Based on generation conditions (access time, token counts, and other data that do not depend on the privacy of the input data), perform statistical processing and machine learning to update the minimum reserved memory amount from step 1.

* When acceptance of requests is stopped, in a commercial environment the requests are routed to another LLM serving node. When all serving nodes are exhausted, requests go into a queue. When the queue is also full, the result is too many requests. In other words, the challenge of this effort is how to avoid too many requests with as few serving nodes as possible, and that is where we get to show our skills.

How the Proposed Method Solves the Batching Problem

The proposed method also mentions moving from per-request processing to per-iteration processing. The concrete approach to that solution lies in GPU kernel processing, so we will skip the explanation of the special kernel processing in this article, but we would like to write another article on how much throughput improves.

Summary

We have briefly introduced and explained the ideas behind PagedAttention.

New LLM serving technologies are being developed every day, by us and by others, and the pace of evolution is fast. What makes PagedAttention, introduced here, especially significant is that it has been released not just as a paper and sample code, but as vLLM, a high-performance implementation that is already starting to be adopted in commercial environments.

We also gave a brief introduction to our own LLM serving technology in this post. Building on the LLM serving technology that PagedAttention has advanced, we plan to connect it to even more efficient GPU operations.

PagedAttention focuses on efficient management of the KV cache, but vLLM goes beyond that and supports many features designed for large-scale LLM serving, such as model parallelism, so we expect its use as an engine to keep spreading.