LLM Internals: KV Cache & Generation

# 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

# 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!

# 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

# 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

messages = client.messages.create(
  model = "claude-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.

# 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

# 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

# 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

# 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