Complete KV Cache and Pruning Discussion Summary

What We Covered

This summary captures our entire discussion about KV cache, pruning strategies, attention mechanisms, and LLM inference.

Part 1: KV Cache Basics

What is KV Cache?

• Purpose: Memory optimization technique to avoid recomputing Key and Value matrices during text generation
• Structure: Matrix format [tokens × channels]
  • Tokens: Sequence positions (words/subwords)
  • Channels: Feature dimensions (model's hidden size, e.g., 768, 1024, 4096)
• Growth: Cache expands as sequence lengthens: [1×channels] → [2×channels] → [3×channels]...

Why KV Cache Matters

• Without cache: Must recompute K,V for all previous tokens when generating each new token
• With cache: Only compute K,V for new token, reuse cached values
• Massive computational savings during inference

Part 2: Pruning Strategies

Core Concepts

• Pruning Direction: Which axis to remove elements from
• Output-Awareness: Using scoring metrics to estimate element importance
• Local Dense Window: Keep recent 32 tokens untouched for accuracy

Per-Channel Pruning

• Method: For each channel (column), selectively remove some token entries
• Direction: Vertical (across tokens)
• Result: Different sparsity patterns for each channel

Per-Token Pruning

• Method: For each token (row), selectively remove some channel entries
• Direction: Horizontal (across channels)
• Result: Different sparsity patterns for each token

Part 3: Output-Awareness Scoring

The Product Formula

Score = Element × Input

Why This Works

• Element: The cached K or V value
• Input: Query value (for K cache) or attention weight (for V cache)
• Product: Captures actual contribution to final output
• Logic: High element + Low input = Low contribution (safe to prune)

Real Example

KV cache element: 0.8
Query input: 0.1
Score: 0.8 × 0.1 = 0.08 (LOW - can prune)

vs.

KV cache element: 0.3
Query input: 0.9  
Score: 0.3 × 0.9 = 0.27 (HIGHER - should keep)

Part 4: Understanding Q, K, V in Attention

What Each Represents

• Query (Q): "What information do I need?" - Search request
• Key (K): "What information do I have?" - Advertisement/label
• Value (V): "Here's the actual information" - Content payload

How They Work Together

1. Q × K: Compute attention weights (who should attend to whom)
2. Softmax: Normalize attention weights
3. Attention × V: Weighted sum of values (what information gets mixed)

Training Perspective

• W_q, W_k, W_v: Three learned transformation matrices
• Same input: Gets transformed three different ways for different purposes
• Learning: Model learns these matrices to solve language modeling task

Example

Token "queen" input: [0.5, 0.8, 0.2, 0.9, ...]

After transformations:
Q = [0.000001, 1, 3, ...] # "I need person/musician info, not royal info"
K = [100, 2, 4, ...]      # "I'm very relevant for royal queries"
V = [1, 2, 3, ...]        # "I contain: royal=1, person=2, musician=3"

Part 5: LLM Inference Process

Two-Phase Approach

Phase 1: Prefilling

• Purpose: Process entire input prompt
• Method: All tokens processed simultaneously (parallel)
• Output: Build initial KV cache, generate first response token
• Speed: Fast due to parallelization

Phase 2: Decoding

• Purpose: Generate response tokens one by one
• Method: Sequential processing, append to KV cache
• Output: Complete response
• Speed: Slower due to sequential nature

Complete Example

Input: "System: You are helpful. User: What is the capital of France?"

Prefilling:
- Process all 17 input tokens at once
- Build KV cache: [17 × hidden_size]
- Generate first token: "The"

Decoding:
Time 1: Add "The" → Cache: [18 × hidden_size] → Generate "capital"
Time 2: Add "capital" → Cache: [19 × hidden_size] → Generate "of"
Time 3: Add "of" → Cache: [20 × hidden_size] → Generate "France"
...continue until complete response

Key Insights Gained

1. KV Cache is Essential: Enables efficient autoregressive generation
2. Pruning is Nuanced: Different strategies (per-channel vs per-token) serve different purposes
3. Output-Awareness is Smart: Considers both stored information and current needs
4. Q,K,V Have Distinct Roles: Not just different values, but different purposes
5. Inference Has Structure: Prefilling vs decoding phases optimize for different constraints
6. Everything Connects: From training objectives to inference efficiency to pruning strategies

Practical Applications

• Memory Optimization: Pruning reduces KV cache size for long sequences
• Inference Acceleration: Smaller cache = faster attention computation
• Quality Preservation: Smart pruning maintains model performance
• Scalability: Enables processing of longer contexts within memory constraints

This comprehensive understanding provides the foundation for working with modern LLM optimization techniques and understanding their trade-offs between efficiency and quality.