Serving LLMs

Summary of notes, taken from Efficiently Serving LLMs from deeplearning.ai

Autoregressive Generation:

How does a Causal Language Model work?

1) Loading the LLM & Tokenizer

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "./models/gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

2) Tokenizing the input prompt

prompt = "The quick brown fox jumped over the"
inputs = tokenizer(prompt, return_tensors="pt")

3) Creating a function that generates the next token

def generate_token(inputs):
  with torch.no_grad():
    outputs = model(**inputs)

  logits = outputs.logits
  last_logits = logits[0, -1, :]
  next_token_id = last_logits.argmax()
  return next_token_id

4) Looper to pass through input tokens and generate 10 next words auto-regressively

generated_tokens = []
next_inputs = inputs
durations_s = []
for _ in range(10):
    t0 = time.time()
    next_token_id = generate_token(next_inputs)
    durations_s += [time.time() - t0]
    
    next_inputs = {
        "input_ids": torch.cat(
            [next_inputs["input_ids"], next_token_id.reshape((1, 1))],
            dim=1),
        "attention_mask": torch.cat(
            [next_inputs["attention_mask"], torch.tensor([[1]])],
            dim=1),
    }
    
    next_token = tokenizer.decode(next_token_id)
    generated_tokens.append(next_token)

print(f"{sum(durations_s)} s")
print(generated_tokens)

KV Cache

To avoid recomputing all the tensors for all tokens at each (word) step, new computed elements are appended to the cache to be used in the next step.

When it is time to generate a new token, the model calculates the attention for the new token. The new K, V pair is appended to the existing K, V tensors that are stored in the cache.

1) Prefill

• The model stores the K, V vectors for all the input prompt into the cache ("The cat sat on"). The next word is auto-regressively generated ("mat")

2) Cached Layer

• K, V are cached, Q is generated for the next new word

3) Decoding

• Only the single previously generated token is fed into the model ("mat").
• The model calculates the K, V vectors for ONLY this token, then retrieves previous K,V from the cache ("The cat sat on"), append new K, V vectors to the cache
• Attention is computed using this updated K,V to generate the next token ("eos") [end of sentence].