LLM Internals: KV Cache & Generation

WITHOUT KV-CACHE: quadratic hell

Generating the t₅ token requires attention throughout the entire sequence:
  Q = W_Q · [t₀, t₁, t₂, t₃, t₄, t₅] ← 6 tokens
  K = W_K · [t₀, t₁, t₂, t₃, t₄, t₅] ← 6 tokens
  V = W_V · [t₀, t₁, t₂, t₃, t₄, t₅] ← 6 tokens
                                           ↑ We are counting everything!

Total for 200 tokens: 1+2+3+...+200 = 20100 matmul steps
Complexity: O(n² × d_model) ← disaster on long prompts

───────────────────────────────── ─────────────────────────────────
With KV-CACHE: linear decode

Key observation: K and V for t₀..t₄ do not change!
  Why recount what has already been counted?

kv_cache = {
  K: [k₀, k₁, k₂, k₃, k₄], ← calculated once, stored
  V: [v₀, v₁, v₂, v₃, v₄], ← calculated once, stored
}

Generation t₅:
  q₅ = W_Q · t₅ ← new token only
  k₅ = W_K · t₅ ← new token only
  v₅ = W_V · t₅ ← new token only
  kv_cache.K.append(k₅) ← update the cache
  kv_cache.V.append(v₅)
  out = attention(q₅, kv_cache.K, kv_cache.V) ← O(seq) per step

Total for 200 tokens: 200 matmul steps ← O(n) at the decode phase
Prefill (prompt processing) is still O(n²) - but only once

# Cache form for one model:
kv_cache.shape = [
  2, #K and V (two tensors)
  num_layers, #32 for Llama-3 8B
  num_kv_heads, #8 (GQA) or 32 (MHA)
  seq_len, # current length (growing!)
  d_head #128 = 4096 / 32
]

# Q is NOT cached because:
# Q is only needed to calculate the current step
# q_i = W_Q t_i - only for new token i
# K and V are ALWAYS needed with each new token

# At each decode step:
kv_cache.K[layer][head].append(k_new)
kv_cache.V[layer][head].append(v_new)
# seq_len increases by 1 → cache grows

WHY IS IT MATHEMATICALLY CORRECT TO CACH K,V?

Causal attention matrix (lower triangle):

       t₀ t₁ t₂ t₃
  t₀ [α₀₀ 0 0 0 ]
  t₁ [α₁₀ α₁₁ 0 0 ]
  t₂ [α₂₀ α₂₁ α₂₂ 0 ]
  t₃ [α₃₀ α₃₁ α₃₂ α₃₃]

When calculating t₃, k₀,k₁,k₂ and v₀,v₁,v₂ are used
These values will not change in subsequent steps:
  k₄ will not be added to k₀, k₀ will remain k₀

→ It is safe to cache K,V: they are immutable after writing
→ Only K and V of the new step are added to the end of the cache
→ Q of the new step is calculated on the fly and is not needed later

# KV-cache size formula (fp16):
kv_bytes = 2 × L × H_kv × S × d_h × 2

# Llama-3 8B: L=32, H_kv=8 (GQA), d_h=128
S= 4,096: 2 × 32 × 8 × 4096 × 128 × 2 = 537 MB
S= 32,768: 2 × 32 × 8 × 32768 × 128 × 2 = 4.3 GB
S=128,000: 2 × 32 × 8 × 128,000 × 128 × 2= 16.8 GB
                                     ↑ MORE models (8GB)

# GPT-4 levels (96 layers, 128 heads, d_h=128, MHA):
S=128,000: 2 × 96 × 128 × 128,000 × 128 × 2= ~640 GB
          ← requires several A100 only for KV-cache!

PREFILL PHASE

  Input: entire prompt [t₀, t₁, ..., t_N] at once
  What it does: simultaneous forward pass for N tokens
              builds KV-cache for all positions
  Bottleneck: COMPUTE (matmul) ← GPU is 100% busy
  Metric: TTFT (Time-To-First-Token)
  Complexity: O(N²) — quadratic of prompt length

  "mom washed the frame" (4 tokens) → prefill 1 forward pass
  long RAG prompt (2000 tokens) → 500000× more expensive

───────────────────────────────── ─────────────────────────────────
DECODE PHASE

  Input: ONE token per step (last generated)
  What it does: forward for 1 token + reads the entire KV-cache
              adds k_new, v_new to cache
  Bottleneck: MEMORY BANDWIDTH ← reading KV-cache from HBM
  Metric: TPS / TPOT (Tokens Per Second, Time Per Output Token)
  Complexity: O(S) per step, O(N×S) total

  GPU utilization during decode: 5–20% ← GPU is empty!
  All the time is busy reading KV-cache from memory

───────────────────────────────── ─────────────────────────────────
LATENCY METRICS

  TTFT = Time To First Token ← prefill latency
  TPOT = Time Per Output Token ← decode speed
  TPS = Tokens Per Second ← 1/TPOT
  E2E = TTFT + N_output × TPOT ← total latency

  Streaming UX: TTFT is more important than E2E
  Batch offline: more important than TPS

# Without chunked prefill:
User A prefill: [────────────────── 2000 tokens ──]
User B decode: [waiting]
                  ← B is blocked for the entire time prefill A

# With chunked prefill (chunk=256):
User A chunk 1: [──256──]
User B token 1: [─1─]
User A chunk 2: [──256──]
User B token 2: [─1─]
...
← A finishes a little later, but B doesn't freeze
← latency tail (p99) improves significantly

# Implementation: vLLM 0.4+, TensorRT-LLM

MHA: Multi-Head Attention (GPT-3, original transformer)

  num_heads_Q = 32
  num_heads_KV = 32 ← K,V for each Q-head

  Q-heads: [Q₀ Q₁ Q₂ ... Q₃₁]
  K-heads: [K₀ K₁ K₂ ... K₃₁] ← 32 separate Ks
  V-heads: [V₀ V₁ V₂ ... V₃₁] ← 32 individual Vs

  KV-cache size: 2 × 32 × S × d_h × 2 bytes ← basic

───────────────────────────────── ─────────────────────────────────
GQA: Grouped Query Attention (Llama-3, Mistral, Gemma)

  num_heads_Q = 32 ← same number of Q-heads
  num_heads_KV = 8 ← a group of 4 Q-heads shares one KV-head

  Q-groups: [Q₀ Q₁ Q₂ Q₃] [Q₄ Q₅ Q₆ Q₇] ... [Q₂₈..Q₃₁]
  K-heads: [K₀] [K₁] ... [K₇]
  V-heads: [V₀] [V₁] ... [V₇]

  KV-cache size: 2 × 8 × S × d_h × 2 bytes ← 4× smaller!
  Quality drop: minimal (GQA ≈ MHA on most benchmarks)

───────────────────────────────── ─────────────────────────────────
MQA: Multi-Query Attention (PaLM, Falcon)

  num_heads_Q = 32 ← same number of Q-heads
  num_heads_KV = 1 ← ALL Q-heads share ONE K,V pair

  KV-cache size: 2 × 1 × S × d_h × 2 bytes ← 32× less!
  Quality drop: noticeable - GQA is usually a better compromise

# Traditional KV-cache:
kv[req_1] = allocate(max_seq_len=2048) ← reserve everything at once
kv[req_2] = allocate(max_seq_len=2048) ← 2048 more, although 100 are needed
← 90% of GPU memory is fragmented with empty blocks

# PagedAttention:
page_size = 16 # tokens per page

req_1 ← [page_4, page_7, page_12] ← not physically adjacent
req_2 ← [page_1, page_9]
req_3 ← [page_2, page_5, page_11] ← allocated as they grow

# block_table stores the mapping of logical → physical pages
# GPU core vLLM can perform attention on non-adjacent blocks

Result:
  memory waste: 30–40% (naive) → < 4% (PagedAttention)
  throughput: +2–4× with the same GPU

Full KV-cache:
  attention(q_i, K[0..i], V[0..i]) ← all tokens
  maximum quality
  memory: O(seq_len) → grows without limit

Sliding Window (W=4096):
  attention(q_i, K[i-W..i], V[i-W..i]) ← window only
  quality: good for local dependencies
  memory: O(W) → constant after W tokens!

SWA problem:
  token outside window = "forgotten"
  distant dependencies (beginning → end of text) are lost

Hybrid: some layers full attention,
        part of the SWA (Mistral Mixtral) layers

WITHOUT PROMPT CACHING: every time again

Request 1: [system_prompt=2000 tok] + [query_1=20 tok]
  → prefill 2020 tokens completely $$$

Request 2: [system_prompt=2000 tok] + [query_2=15 tok]
  → prefill 2015 tokens completely $$$
                                     ← 2000 tok is being recalculated!

───────────────────────────────── ─────────────────────────────────
WITH PROMPT CACHING: KV-cache on the server

Request 1: [system_prompt=2000 tok] + [query_1=20 tok]
  → prefill 2020 tokens cache_write
  → KV-cache for system_prompt is saved on the server

Request 2: [system_prompt=2000 tok] + [query_2=15 tok]
  → cache_hit: 2000 tokens from cache 10% cost
  → prefill only 15 new tokens → TTFT is 10x faster

───────────────────────────────── ─────────────────────────────────
KEY RULE OF USE

  Static → start of prompt (cached)
  Dynamic → end of prompt (not cached)

  ✓ Correct:
    [system_prompt] ← static → cached
    [rag_documents] ← static for TTL period → cached
    ────────────────
    [user_query] ← changes every time

  ✗ Incorrect:
    [user_query] ← first
    [system_prompt] ← after → is not cached!

messages = client.messages.create(
  model = "cloud-sonnet-4-6",
  system = [
    {
      "type": "text",
      "text": system_prompt, ← static prompt
      "cache_control": {"type": "ephemeral"}
    }
  ],
  messages = [
    {"role": "user", "content": [
      {"type": "text", "text": rag_context,
       "cache_control": {"type": "ephemeral"}},
      {"type": "text", "text": user_query} ← not cacheable
    ]}
  ]
)

# Check the result:
usage = messages.usage
usage.cache_read_input_tokens # how many from cache
usage.cache_creation_input_tokens # how many are written to the cache
usage.input_tokens # regular (not cached)

OpenAI - automatically:
  ✓ No cache_control markup needed
  ✓ The API itself determines matching prefixes
  ✓ 50% discount on cached tokens (input)
  ? There is no explicit control - you don’t know what is cached

# Check in response:
response.usage.prompt_tokens_details
  → {"cached_tokens": 1920, "audio_tokens": 0}

Anthropic - explicit markup:
  ✓ Full control over which block to cache
  ✓ 90% discount (vs 50% for OpenAI) → more aggressive
  ! We need to rebuild the prompt structure
  ! Minimum 1024 tokens per block

Advice: both providers - install static first,
dynamic (query) to the end. It works for both.

STEP 1: RAW LOGITS (after LM Head)

  logits = W_lm_head hidden_state[-1]
  shape: [128 256] ← vocab_size Llama-3

  Top 5 example:
  "cat" : 8.41
  "dog" : 6.23
  "mouse" : 5.87
  "bird" : 4.12
  ...128251 other tokens: -12 to +3

───────────────────────────────── ─────────────────────────────────
STEP 2: TEMPERATURE (scaling)

  logits = logits / T

  T=0.0: logits/→∞ → argmax (greedy, deterministic)
  T=0.7: “cat”: 12.01, “dog”: 8.90 ← the gap has increased
  T=1.0: "cat": 8.41, "dog": 6.23 ← raw logits
  T=1.5: “cat”: 5.61, “dog”: 4.15 ← the gap has shrunk, chaos

───────────────────────────────── ─────────────────────────────────
STEP 3: TOP-K (tail trimming by number)

  k=50: keep only 50 tokens with the highest logit
  the rest 128206 → logit = −∞

───────────────────────────────── ─────────────────────────────────
STEP 4: TOP-P / NUCLEUS SAMPLING (trimming by total probability)

  After softmax we accumulate probabilities:
  "cat": 0.62 → cumsum = 0.62
  "dog": 0.18 → cumsum = 0.80
  "mouse": 0.09 → cumsum = 0.89
  "bird": 0.04 → cumsum = 0.93 ← if p=0.9, cut off here
  ...the rest → −∞

  p=0.9: keep tokens until Σprob ≤ 0.9

───────────────────────────────── ─────────────────────────────────
STEP 5: SOFTMAX + SAMPLE

  probes = softmax(filtered_logits)
  next_token = random.choice(tokens, p=probs)
  ← the only non-deterministic step!

# Initial logits: [8.41, 6.23, 5.87, 4.12]

T=0.3 (cold, “confident”):
  scaled: [28.0, 20.8, 19.6, 13.7]
  softmax: [0.993, 0.004, 0.002, 0.001]
  ← almost always "cat"

T=1.0 (neutral, “raw distribution”):
  scaled: [8.41, 6.23, 5.87, 4.12]
  softmax: [0.617, 0.083, 0.058, 0.011]
  ← sometimes other options

T=2.0 (hot, “chaotic”):
  scaled: [4.21, 3.12, 2.93, 2.06]
  softmax: [0.342, 0.189, 0.158, 0.071]
  ← often unexpected tokens

# T→0: argmax (deterministic)
# T→∞: uniform (everything is equally likely)

Repetition penalty (θ = 1.3):
  if token_id in context:
    logit[token_id] = logit[token_id] / 1.3 ← if > 0
    logit[token_id] = logit[token_id] * 1.3 ← if < 0
  # 1.0 = no penalty, 1.3 = moderate, 2.0 = aggressive

Logit bias:
  # force deny token
  logit_bias = {token_id_of_swear: -100} ← = -∞

  # force the token to be strengthened
  logit_bias = {token_id_of_yes: +10} ← ≈ 20000× more likely

# OpenAI API: logit_bias parameter
# Usage: guardrail for banned words
# or force structured output

MECHANICS

Draft model (haiku/Llama-3-1B):
  Generates K=5 tokens quickly:
  ["neural", "network", "it", "mathematical", "function"]
  Cost: 5 × 1B-forward-pass ← cheap

Target model (cloud-sonnet / Llama-3-70B):
  One forward pass by context + 5 draft tokens:
  Receives 5 probability distributions simultaneously
  Cost: 1 × 70B-forward-pass ← but validates 5 tokens!

VERIFICATION (rejection sampling):

  Draft: ["neural", "network", "it", "mathematical", "function"]
  Target: P(neural)=0.71 P(network|...)=0.83 P(this|...)=0.04

  Token 1 "neural": draft P=0.7, target P=0.71 → ACCEPT
  Token 2 "network": draft P=0.8, target P=0.83 → ACCEPT
  Token 3 "this": draft P=0.6, target P=0.04 → REJECT
  → regenerate from position 3, discard the rest

RESULT

  Traditional decode: 5 steps × 1 forward = 5 × 70B passes
  Speculative decode: 5 draft + 1 verify = 5 × 1B + 1 × 70B
                                             ≈ 2–3× faster with high acceptance rate

  Acceptance rate ≈ 70–80% with draft~target in the family
  Self-speculation: the model itself generates a draft based on early layers

# Size of Llama-3 70B in different formats:
fp32: 280 GB ← unrealistic
fp16: 140 GB ← training
bf16: 140 GB ← training (better fp16 for large numbers)
int8: 70 GB ← almost no quality loss (LLM.int8)
int4: 35 GB ← GPTQ/AWQ: slight loss of quality
int3: 27 GB ← noticeable loss
int2: 18 GB ← significant loss

# KV-cache can also be quantized:
fp16 → int8: KV-cache 2x smaller
  ← saves GPU memory for larger batches
  ← slight drift on long sequences

# Standard attention (slow, lots of memory):
S = Q @ K.T / sqrt(d_k) # [n × n] in HBM - expensive!
A = softmax(S) # [n × n] in HBM
O=A@V#[n×d] in HBM
Memory: O(n²)

# Flash Attention (Dao et al. 2022):
# Tile algorithm: Q, K, V blocks are loaded into SRAM
# (fast GPU memory), softmax is calculated incrementally
#S never fully materializes
Memory: O(n) ← linear! only SRAM tiles

Result:
  Speed: 2–4x faster on A100
  Memory: 5–20× less
  Numerically identical to standard attention

TOKEN EMBEDDINGS (within LLM)

  What: matrix E ∈ ℝ^[vocab_size × d_model]
          lookup table: token_id → vector

  Input: integer token-ID (15496)
  Output: row vector of matrix E, dim = d_model (4096)
  Meaning: initial representation of a token in model space
          changes (enriches) on each transformer block

  Example:
  "mother" → id=77341 → E[77341] → [0.12, -0.34, 0.89, ...] ∈ ℝ^4096

  Usage: only inside the model, first layer
  NOT SUITABLE for search: not trained for retrieval task,
  context-sensitive (contextualized through attention)

───────────────────────────────── ─────────────────────────────────
TEXT EMBEDDINGS (for RAG/search)

  What: output of a specialized encoder model
          entire text → one vector representation

  Input: text of arbitrary length (up to max_seq_len)
  Output: one vector, dim = 768–3072 (depending on model)
  Meaning: semantic representation of the ENTIRE text
          specially trained: close in meaning → close vectors

  Example:
  "mom washed the frame" → encoder → [0.41, 0.23, -0.18, ...] ∈ ℝ^768
  cosine("mom washed the frame", "mom washed the frame") → 0.97 ← close!

  Usage: indexing + search in RAG
  Models: text-embedding-3-large, bge-m3, e5-large

───────────────────────────────── ─────────────────────────────────
KEY DIFFERENCE

  Token embedding: "bank" → always the same vector ← static
  Text embedding: "bank on the river" → one vector ← context included
                    "bank account" → another vector
  LLM hidden state: "bank" → different vector in different contexts ← contextualized

# E5-mistral: prompt engineering for retrieval
prompt = f"Instruct: Given a query, find relevant documents.\n
Query: {query}"
embedding = llm_model.last_hidden_state(prompt)[-1]
← take the last hidden state (pool)

# When justified:
✓ the task requires deep understanding (code, science)
✓ has a GPU, latency < 1s is acceptable
✓ bi-encoder embedding is not enough

# Why standard text-embedding is better in 99% of cases:
✗ LLM inference: 100–500ms vs 5–20ms for bi-encoder
✗ cost 10–50× higher
✗ bge-m3 on MTEB is often no worse than LLM-embedded

INPUT: "Continue the text: mom washed the frame"
────────────────────────────── ───────────────────────────────

① TOKENIZATION (tokenizer: tiktoken / sentencepiece)
  "Continue" → [3643, 70878] ← 2 tokens (Cyrillic is expensive)
  "text" → [50078]
  "mom" → [86392]
  "soap" → [42341, 261] ← the word broke
  "frame" → [100342]

  input_ids = [3643, 70878, 50078, 86392, 42341, 261, 100342]
  seq_len = 7

────────────────────────────── ───────────────────────────────
② TOKEN EMBEDDINGS (lookup in matrix E ∈ ℝ^[128256 × 4096])
  E[3643] → x₀ ∈ ℝ^4096 ← vector for "Continue"
  E[70878] → x₁ ∈ ℝ^4096
  ...
  E[100342]→ x₆ ∈ ℝ^4096
  shape: [7 × 4096]

────────────────────────────── ───────────────────────────────
③ POSITIONAL EMBEDDINGS (RoPE)
  # RoPE is not applied as addition, but through rotation Q,K
  # when calculating attention inside each block
  # position is encoded into sine/cosine frequencies
  position_ids = [0, 1, 2, 3, 4, 5, 6]

────────────────────────────── ───────────────────────────────
④ 32 TRANSFORMER BLOCKS (repeat 32 times)

  ─ Block 0: ──────────────────────────────────────
  RMSNorm(x) → normalization: x / rms(x) γ

  GQA (32Q/8KV heads):
    Q = x W_Q shape: [7 × 32 × 128]
    K = x · W_K shape: [7 × 8 × 128] ← GQA: only 8!
    V = x W_V shape: [7 × 8 × 128]

    # K and V are placed in kv_cache[block_0]:
    kv_cache[0].K = K ← 7 tokens × 8 heads × 128 = 7168 float16
    kv_cache[0].V = V

    # Causal attention (lower triangle only):
    scores = Q Kᵀ / √128 shape: [7 × 32 × 7]
    scores += causal_mask ← +0 or −∞
    α = softmax(scores)
    attn_out = α V shape: [7 × 32 × 128]
    attn_out = reshape + W_O → [7 × 4096]

  Residual: x = x + attn_out

  RMSNorm(x)

  SwiGLU FFN: 4096 → 14336 → 4096
    gate = x W_gate [7 × 14336]
    up = x W_up [7 × 14336]
    act = SiLU(gate) * up
    down = act W_down [7 × 4096]

  Residual: x = x + down
  ─ ... repeated 31 times ───────────────────────

  Block output: x₃₂ ∈ ℝ^[7 × 4096]

────────────────────────────── ───────────────────────────────
⑤ LM HEAD (projection onto dictionary)
  x₃₂[-1] ← take only the last position “frame”
  logits = x₃₂[-1] · W_lm_headᵀ shape: [128256]

  Top 5 logits:
  "erased": 9.21 ← likely successor!
  "rinsed": 7.84
  "washed": 6.91
  "cleaned": 5.43
  "rubs": 4.12

────────────────────────────── ───────────────────────────────
⑥ SAMPLING (T=0.7, top_p=0.9)
  logits / 0.7 → [13.16, 11.20, 9.87, 7.76, 5.89, ...]
  softmax → [0.71, 0.12, 0.05, 0.02, 0.01, ...]
  top_p=0.9 → leave Σprob ≤ 0.9 for now
  sample() → "erased" (id: 91234)

────────────────────────────── ───────────────────────────────
⑦ NEXT STEP (DECODE)
  input_ids = [...7 tokens..., 91234] ← add a new one
  kv_cache contains K,V for the first 7 positions
  Compute only one new token: q₇, k₇, v₇
  kv_cache.append(k₇, v₇)
  ← decode step: O(1) mathmul, O(seq) cache read

# Llama-3 8B, fp16, A100 80GB SXM, vLLM

Prefill (compute-bound):
  1024 prompt tokens: TTFT ≈ 30 ms
  4096 prompt tokens: TTFT ≈ 120 ms
  16K prompt token: TTFT ≈ 900 ms ← O(n²) hits

Decode (memory-bound):
  batch=1: TPS ≈ 80 tok/s
  batch=8: TPS ≈ 220 tok/s ← batching helps
  batch=32: TPS ≈ 400 tok/s ← saturates HBM bandwidth

Memory:
  Model weights: ~16 GB (fp16)
  KV-cache @ 4K × batch=8: ~4.3 GB
  Activations: ~1–2 GB
  Total A100 80GB: fits with a margin

With optimizations:
  int4 (AWQ): 4 GB model, batch×2
  Speculative: latency 2–3× lower
  Prompt cache: TTFT 5–10× lower on hit