Careeroza

Real World Example

Where This Knowledge Pays Off

  • Self-hosting open models (vLLM, TGI, Ollama): GPU memory goes to weights and the KV cache, which grows with context length and concurrent users.
  • Search and recommendations: encoder models (BERT family) power semantic search at Google, Flipkart and others.
  • Cost and latency: understanding n² attention and the KV cache explains why long prompts, long chat histories and long outputs are slow and costly with any provider.

KV Cache Memory for a Serving Plan

Example

ShopEase wants to self-host Llama 3.1 8B for an internal tool. How much GPU memory does the KV cache need? Per token: 2 (K and V) × layers × key/value heads × head size × bytes.

layers, kv_heads, head_dim, bytes_fp16 = 32, 8, 128, 2
per_token = 2 * layers * kv_heads * head_dim * bytes_fp16
print(f"KV cache per token: {per_token // 1024} KB")

weights_gb = 8e9 * 2 / 1e9
for context, users in [(4_000, 1), (8_000, 10), (32_000, 10)]:
    kv_gb = per_token * context * users / 1e9
    print(f"{users:>2} users x {context:>6,} tokens: KV {kv_gb:5.1f} GB + weights {weights_gb:.0f} GB = {kv_gb + weights_gb:5.1f} GB")

Output

KV cache per token: 128 KB
 1 users x  4,000 tokens: KV   0.5 GB + weights 16 GB =  16.5 GB
10 users x  8,000 tokens: KV  10.5 GB + weights 16 GB =  26.5 GB
10 users x 32,000 tokens: KV  41.9 GB + weights 16 GB =  57.9 GB

Looking Inside a Real Model

from transformers import AutoTokenizer, AutoModelForCausalLM

name = "Qwen/Qwen2.5-0.5B-Instruct"                 # small enough for a laptop CPU
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name)

inputs = tok("ShopEase refunds take", return_tensors="pt")
out = model(**inputs, output_attentions=True)
print("layers:", len(out.attentions), "| heads:", out.attentions[0].shape[1])
next_id = out.logits[0, -1].argmax().item()
print("next token:", repr(tok.decode(next_id)))

Sample output

python inspect_model.py

Example output

layers: 24 | heads: 14
next token: ' '