1

Inference Systems
What is inference?

3
 forward vs generate
● Forward: single forward pass ● Generate:
through LLM ○ Generates multiple tokens,
until EOS or max tokens
reached

4
source : [Link]
 What is inference?

● Text generation (LM Head)

● Classification
● Translation

[Link] 5
ng-a-gpt-style-llm-classifier
x x x time
1 2 3
c c c3
1 2
Layer 2
MLP MLP MLP
attn attn attn
v v v
k1 1 k2 2 k3 3
q q q
1 2 3

b1 b2 b3

MLP MLP MLP Layer 1

attn attn attn

v v v
k1 1 k2 2 k3 3
q q q
1 2 3

a1 Once a2 upon a3 a 6
 Simplest generate example

from transformers import AutoTokenizer, AutoModelForCausalLM

import torch

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

prompt = "My trip to Yosemite was"

inputs = tokenizer(prompt, return_tensors="pt")

output = [Link](inputs.input_ids, max_length=100)

generated_text = tokenizer.batch_decode(output, skip_special_tokens=True)

print(generated_text)

7
Inference engines

8
What are inference engines?
● Wrap systems concepts around the process
of inference
○ Batching
○ Streaming
○ Fault management
○ Error handling
○ Parallelism (multi-GPUs)
● Optimizations:
○ Memory storage optimizations
○ Offloading
● Most popular open-source is vllm. Others
include SGLang, NVIDIA’s Tensor-RT LLM 9
[Link]
A Survey on Inference
Engines for Large
Language Models:
Perspectives on
Optimization and
Efficiency
([Link]
05.01658)10
Terminology

11
Prefill vs decode
● Prefill: input/prompt phase. Generation of the first output token (ingestion of the
entire input in parallel)

12
Prefill vs decode
● Prefill: input/prompt phase. Generation of the first output token (ingestion of the
entire input in parallel)
● Decode: output/generation phase. Decoding of subsequent output tokens
(sequential)

13
Prefill time Decode
c2 c3
Layer 2 MLP MLP MLP MLP Layer 2
attn attn attn attn
v v v
k1 1 k2 2 k3 3 v
q q q k4 4
q
1 2 3
4

b1 b2 b3 b4

Layer 1 MLP MLP MLP MLP Layer 1

attn attn attn attn

v v v v
k1 1 k2 2 k3 3 k4 4
q q q q
1 2 3 4

Input: a a2 a3 a4
Once upon a time
“Once upon a” 1 14
 time , Decode
Prefill
c2 c3
Layer 2 MLP MLP MLP MLP
attn attn attn attn
v v v
k1 1 k2 2 k3 3 v
q q q k4 4
q
1 2 3
4

b1 b2 b3 b4

Layer 1 MLP MLP MLP MLP

attn attn attn attn

v v v v
k1 1 k2 2 k3 3 k4 4
q q q q
1 2 3 4

Input: a a2 a3 a4
Once upon a time
“Once upon a” 1 15
Metrics

● Throughput: Number of Tokens processed per second

● Latency:
○ Time to first token (prefill time)
○ Time between tokens (decode time)

Time Between Tokens

Time To First Token 16
(inter-token latency)
Let’s write our own generate

18
 Make sure pytorch_model.bin is in the exercises directory
wget [Link]
Step-wise plan - pytorch_model.bin
- simple
- kvcache

First let’s take a look at [Link]

1. Create position ids tensor (What is the position ids for this input?) [1, 2, … n]
2. Call forward(inputs, position_ids) (Make sure both are tensors)
3. What is output (logits) shape? We will learn about this later
4. Append output token to the inputs tensor
5. Wrap entire code in a “for” loop (What is the exit condition?)

19
MAX_SEQ_LEN = 10
config = GPT2Config.from_pretrained('gpt2')
model = GPTLMHead(config)
tokenizer = AutoTokenizer.from_pretrained("gpt2")

def generate(input):
# Tokenize input
tokenized = [Link](input)
inputs = [Link](tokenized)
See the signature of
while ???: # What is exit condition of the loop? [Link]()
positions = ??? # Create position ids tensor
logits = model(???) # Call [Link] with inputs and position ids
logits = logits[-1, :]
next_token = [Link](logits, dim=-1, keepdim=True)
inputs = ??? # Concatenate next_token to inputs for next step here

output = [Link](inputs)
return output

input = 'The quick brown fox jumped' 20

output = generate(input)
MAX_SEQ_LEN = 10
config = GPT2Config.from_pretrained('openai-community/gpt2-large')
model = GPTLMHead(config)
tokenizer = AutoTokenizer.from_pretrained("gpt2")

def generate(input):
# Tokenize input
tokenized = [Link](input)
inputs = [Link](tokenized,)

while len(inputs[0]) < MAX_SEQ_LEN:

positions = [Link](len(inputs))
logits = model(inputs, positions)
logits = logits[-1, :]
next_token = [Link](logits, dim=-1, keepdim=True)
inputs = [Link]((inputs, next_token), dim=0)

output = [Link]([Link]())
return output

input = 'The quick brown fox jumped' 21

output = generate(input)
Core features

22
Core features of inference engines
Inference engines such as vllm etc. have certain features/optimizations

We will look at 3 such techniques:

1. KV Caching
2. Continuous Batching
3. Paged Attention

23
KV Caching

24
Why cache KVs?
What are K and V?
For causal self-attention, K/V of all past tokens is required to calculate
attention of current token
Problem: It is recalculated again and again for every token
Idea: Generate and cache once, reuse for subsequent tokens

25
Recap of
self attention

26
Q.K is a score:
how relevant is a past
token (through its K
representation) to the
current token
(through its query
representation)

27
Finally, each past token (through its Value
Representation) is multiplied with the score and
summed up to obtain a new representation of the
current token

28
 What does KV recomputation mean?

v v v
k1 1 k2 2 k3 3
q q q
1 2 3

b1 b2 b3

1. We need the K and V

of all tokens in the past,
v1 at every layer
k1 v
k2 2
v
k3 3 2. This requires
q1 q q calculation of b1, b2 etc.
2 3 for all past tokens
a1 Once a2 upon a3 a
29
Let’s see the repetition in action
1. GPTAttention has a field [Link] which refers to the layer
2. Now print the shape of the k or v tensor for a particular layer after the first 5
lines of the attention block before SDPA

Head size

[Link]([12, 5, 64])

Num heads

30
 Num heads Head size
Let’s see the repetition in action
[Link]([12, 5, 64])
1. GPTAttention has a field [Link] which refers to the layer
2. Now print the shape of the k or v tensor for a particular layer
3. One dimension grows every iteration - that is the number of tokens
4. Pick a random idx and print keys[a][b][c] (Choose random valid values of a, b,
c)
class GPTAttention([Link]):
def forward(self, x):
batch_size, seq_len, _ = [Link]
q, k, v = …
queries = …
keys = …
values = …

if [Link] == 0:
print(keys[7][4][29])
31
 x2 = “a”

c2

Layer 2

v v
k1 1 k2 2 q1 = Q.a1
q q
b1is output after layer 1
1 2 For every token, only its q matters
Causal attention: For every token, only its
b2 and past tokens k,v matter
b1

Layer 1

v v
k1 1 k2 2
q q
1 2
k1= K.a1
a1 Once a2 upon
33
 c3

Layer 2

v v v
k1 1 k2 2 k3 3
q q q
1 2 3

b1 b2 b3

Layer 1

v v v
k1 1 k2 2 k3 3
q q q
1 2 3
k1= K.a1
a1 Once a2 upon a3 a
34
 c3

v v v
k1 1 k2 2 k3 3
q q q
1 2 3

b1 b2 b3

v v v
k1 1 k2 2 k3 3
q q q
1 2 3
k1= K.a1
a1 Once a2 upon a3 a
35
 c3

v v v
k1 1 k2 2 k3 3
q q q
1 2 3

b1 b2 b3

We don’t pass entire

input for every
decode run, only the
v v v
k1 1 k2 2 k3 3 last token
q q q
1 2 3
k1= K.a1
a1 Once a2 upon a3 a
36
 c3
What is the
size of the
v
k1 1
v
k2 2 v
k3 3 KV cache?
q q q
1 2 3

b2 b3 How many KVs are we

b1
storing?
Num layers x
Num tokens x
2
v v v
k1 1 k2 2 k3 3
q q q
1
k1= K.a1
2 3
Size of K/V?:
Embedding size
a1 Once a2 upon a3 a
37
 Only 12 lines to be added
Let’s implement KV Cache ([Link]) Only 9 lines to be changed

1. Create an argument called kv_cache to the forward in all the

modules calling Attention
a. GPTLMHead, GPTModel, GPTBlock, GPTAttention

38
 Only 12 lines to be added
Let’s implement KV Cache ([Link]) Only 9 lines to be changed

1. Create an argument called kv_cache to the forward in all the

modules calling Attention
a. GPTLMHead, GPTModel, GPTBlock, GPTAttention

class GPTModel([Link]):
def forward(self, x, kv_cache=None):
…
class GPTLMHead([Link]):
def forward(self, inputs, position_ids, kv_cache=None):
… 39
 Only 12 lines to be added
Let’s implement KV Cache ([Link]) Only 9 lines to be changed

1. Create an argument called kv_cache to the forward in all the

modules calling Attention
b. In GPTModel, kv_cache should be initialized as [None] x n_layers
for each block. Now, we’ll pass the kv_cache[i] to each block’s
forward call

40
 Only 12 lines to be added
Let’s implement KV Cache ([Link]) Only 9 lines to be changed

1. Create an argument called kv_cache to the forward in all the

modules calling Attention
b. In GPTModel, kv_cache should be initialized as [None] x n_layers
for each block. Now, we’ll pass the kv_cache[i] to each block’s
forward call
class CausalModel([Link]):
def forward(self, inputs, position_ids, kv_cache=None):
if kv_cache is None:
kv_cache = [None for _ in range(len(self.h))]
…
h(x, kv_cache[i]) 41
Let’s implement KV Cache ([Link]:Attention)

2. Init case (Prefill):

i. kv_cache is passed as None
ii. Populate it as tuple of (keys, values)

42
 Let’s implement KV Cache ([Link]:Attention)

2. Init case (Prefill):

i. kv_cache is passed as None
ii. Populate it as tuple of (keys, values)

class GPTAttention([Link]):
def forward(self, kv_cache=None):
…
if kv_cache is None:
kv_cache = (keys, values)
43
Let’s implement KV Cache ([Link]:Attention)

3. Decode case: kv_cache has the tuple (keys, values) of past values
i. Read past_keys (values) from kv_cache. Print shape
ii. Print shape of current keys
iii. See which dimension to concatenate to create the merged
keys and values [Link]((t1, t2), dim)
iv. Continue using the new keys and values tensor

44
 Let’s implement KV Cache ([Link]:Attention)

3. Decode case: kv_cache has the tuple (keys, values) of past values
i. Read past_keys (values) from kv_cache. Print shape
ii. Print shape of current keys
iii. See which dimension to concatenate to create the merged
keys and values [Link]((t1, t2), dim)
iv. Continue using the new keys and values tensor
class GPTAttention([Link]):
def forward(self, kv_cache=None):
if kv_cache is not None:
past_keys, past_values = kv_cache
keys = [Link]((past_keys, keys), dim=2) 45
Let’s implement KV Cache ([Link])

4. Return the populated KV cache at all modules:

a. Return kv_cache at GPTAttention
b. Return at GPTBlock
c. Return at GPTModel (remember each layer will return its own
kv_cache, so this has to be a list)
d. Return at CausalModel back to user

46
Let’s implement KV Cache

5. Changes at generate:
a. Separate the prefill and decode calls. Prefill is once, decode is in a
loop.
b. Input: We will not pass the entire input to decode, just the output
token of the last forward.
c. kv_cache output of one call is passed as input to next call
d. Position ids: What was the old position ids? What should the new
one be?

47
 c3

v v v
k1 1 k2 2 k3 3
q q q
1 2 3

b1 b2 b3

v v v
k1 1 k2 2 k3 3
q q q
1 2 3
k1= K.a1
a1 Once a2 upon a3 a
48
 Attention mask
K
K
Once upon a
Once upon a time
Once

q1.k1 q1.k2 q1.k3

Once
✅ ❌ ❌ ❌
upon

Q q2.k1 q2.k2 q2.k3

Q ✅ ✅ ❌ ❌

upon
q3.k1 q3.k2 q3.k3
a

✅ ✅ ✅ ❌

time a
✅ ✅ ✅ ✅

49
 Attention mask
K
K
Once upon a
Once upon a time
Once

q1.k1 q1.k2 q1.k3

Once
✅ ❌ ❌ ❌
upon

Q q2.k1 q2.k2 q2.k3

Q ✅ ✅ ❌ ❌

upon
q3.k1 q3.k2 q3.k3
a

✅ ✅ ✅ ❌

time a
✅ ✅ ✅ ✅
With KV Cache

50
Attention mask with KV Caching

6. In attn():
a. “is_causal=True” fills the triangular attention mask by default. We want
to turn it off for decode
b. How to know whether prefill or decode? Use [Link] and [Link]
c. If prefill: Use is_causal=True
d. If decode: Use is_causal=False

If everything goes right, your output should match the previous output

51
Savings
Run the timed_generate.py in the solutions file (Copy pytorch_model.bin)
You should see benefits of KV Cache for larger number of tokens

52
Let’s summarize KV Cache
1. Why do we cache KV?
2. What do we save? Compute or memory?
3. How are prefill and decode phases different?
4. What can we do if we run out of memory?

53
Batching

54
Batching

No batching:
● Generate runs 1 request at a time No
● GPU not utilized fully batching

55
Batching

No batching:
● Generate runs 1 request at a time No
● GPU not utilized fully batching
Static batching:
● Request level batching: Static
● Batch a set of requests batching
● More parallelization

56
Static batching

57
 Static Batching
First dimension in inputs is batch size. Previously, we set it to 1.
How to achieve static batching:
1. Batch multiple requests together
2. How to handle different size requests? Padding

What are the

problems in this
approach?

58
 Static batching Entire batch waits until T8

Iterate until the entire batch is over - low util

High latency for finished requests Not utilized fully 60
 Continuous Batching
Iterative/Continuous1 Batching

● Token level batching No

● Replace completed requests with new batching
● Batch size parameter:
○ Throughput vs latency tradeoff Static
batching
Any problems?
Adaptive
batching

1. Orca: A Distributed Serving System for Transformer-Based Generative Models NSDI ’22
2. Image source: [Link] 61
 Next input starts executing
Continuous batching immediately

63
Let’s summarize batching

1. 2 types of batching
2. What do we gain by batching?
3. What do we lose?
a. Multi-step scheduling

64
Paged Attention

65
 Memory management of KV Cache
With batching and KV Caching, the KV cache is initialized as a contiguous Tensor
for the max sequence length supported by the model

For a long context model (8k - 128k tokens), how much memory does this take?

2048 slots reserved for Request 1

R1 Once upon a time …

R2 How to read …

2048 slots reserved for Request 2

66
PagedAttention: [Link]
 Reserved ones are only going to be used External fragmentation: Free space from
later, some other request can use it now past requests

4 KV cache states 4 slots for future 2040 slots never used External
7 slots reserved
(token states) (reserved) (internal fragmentation) fragmentation

Once upon a time … How to read …

Actual KV cache
states stored Internal fragmentation is never used,
so reserved unnecessarily

67
 4 KV cache states 4 slots for future 2040 slots never used External
7 slots reserved
(token states) (reserved) (internal fragmentation) fragmentation

Once upon a time … How to read …

68
 4 KV cache states 4 slots for future 2040 slots never used External
7 slots reserved
(token states) (reserved) (internal fragmentation) fragmentation

Once upon a time … How to read …

69
 4 KV cache states 4 slots for future 2040 slots never used External
7 slots reserved
(token states) (reserved) (internal fragmentation) fragmentation

Once upon a time … How to read …

70
How do they do this? Blocks/Pages

Block - KV Cache data for a few tokens

71
72
 2 - 4x throughput gains

[Link] 73
Decoding strategies

75
 What is decoding?
What is the output of inference? Not a token!
Logits: A degree of similarity to each token in the vocabulary

Output
Input

76
[Link]
 Decoding Strategies
Which token to predict as response?

logits = model(inputs, positions)

print([Link]) [x, 50257]
logits = logits[-1, :]

How to convert these logits to a token?

77
[Link]
 Softmax
Logit (similarity score) -> Probability score (that adds up to 1)
This is done via a softmax normalization function which has a “temperature” parameter
which scales the logits value.

Makes values more extreme Makes values more centered

More More creative

deterministic and diverse 78
Greedy search

Choose word with highest probability

at every step
Problems:
1. It starts repeating itself after
some time
2. Misses high prob. words hidden
behind low prob. words

next_token = [Link](logits, dim=-1, keepdim=True)

79
[Link]
Beam search

● Takes n top probabilities into account

● Runs n decodes in parallel
● Can still get stuck in a loop

How to avoid repetition -> add

randomness

80
Top-K Sampling

● Randomly pick one of the output tokens

● But this could bias towards the long tail
of irrelevant options
● Select only top-K of those and normalize
the probability
● How to decide K?

81
Top-p sampling
K is decided by number of options that add up to a probability p

82
 Let’s see effect of temperature
from transformers import pipeline
prompt = "The quick brown fox jumps over the"
generator(prompt, max_new_tokens=20, do_sample=True, temperature=0.7)
generator(prompt, max_new_tokens=20, do_sample=True, temperature=1.5)
generator(prompt, max_new_tokens=20, do_sample=True,
temperature=0.00001)

83
Research survey

84
Prefix Caching
If prefix is the SAME for multiple requests, they can share the KV Cache.
Techniques to efficiently match prefix and load its KV Cache
- Hash based
- Radix Tree based
Should be fast and not affect the latency in case of non-match
Where do you think this would be useful?

85
 KV Cache size reduction
Problem: KV Cache can grow several times How to reduce KV Cache size: Any ideas?
larger than model size
Sparsification:

● Not all values are equally important

- Use attention mask to determine
which tokens actually “influence”
others

Sliding window attention:

● Only consider last N tokens

86
DejaVu
 Speculative Decoding (Draft model)
Problem: Auto-regressive nature of generation means only 1 token is decoded in one iteration
Reduce latency without compromising output correctness!

87
 Speculative Decoding (Draft model)
Problem: Auto-regressive nature of generation means only 1 token is decoded in one iteration
Reduce latency without compromising output correctness!

Small model:
50 - 200M
params

Step 1: Draft model quickly generates 4-5

tokens (auto-regressively)
88
[Link]
 Speculative Decoding (Draft model)
Problem: Auto-regressive nature of generation means only 1 token is decoded in one iteration
in
Reduce latency without compromising output correctness!

Small model:
50 - 200M
params
(Original model)

Once upon a time

Step 1: Draft model quickly generates 4-5 Step 2: Target model quickly verifies these
tokens (auto-regressively) tokens in parallel. Worst case, re-run from
that point 89
[Link]
 Prefill-Decode Disaggregation
Problem: Prefill is compute intensive, decode is memory intensive

Not optimal to run both with similar parameters / on the same system

90
Graph from DistServe (OSDI’24)
 Prefill-Decode Disaggregation

Solution: Run them on different machines

Overlap communication
with computation
Stream the KV Cache

Several prior art: DistServe, Llumnix, 91

Figure from DistServe (OSDI’24) LMCache, DejaVu