KV Cache, Multi-Query, Grouped-Query Attention

1. The KV Cache Mechanism

During autoregressive generation, the KV cache stores previously computed key-value pairs to avoid redundant computation:

1.1 Standard Attention (No Cache)

For each new token , recompute attention with all previous tokens:

Cost: per token.

1.2 KV Cache

Store and :

For new token :

1. Compute ,
2. Append to cache:
3. Compute attention:

2. KV Cache Memory Layout

<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="400" fill="#f8fafc"/>
  
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">KV Cache Memory Layout</text>
  
  <g transform="translate(50, 60)">
    <text x="350" y="0" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">KV Cache for Layer l, Head h</text>
    
    <g transform="translate(0, 15)">
      <text x="175" y="0" text-anchor="middle" fill="#3b82f6" font-family="monospace" font-size="10" font-weight="bold">Key Cache</text>
      <rect x="0" y="15" width="50" height="40" rx="3" fill="#3b82f6"/>
      <text x="25" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k1</text>
      <rect x="55" y="15" width="50" height="40" rx="3" fill="#3b82f6"/>
      <text x="80" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k2</text>
      <rect x="110" y="15" width="50" height="40" rx="3" fill="#3b82f6"/>
      <text x="135" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k3</text>
      <text x="180" y="40" fill="#64748b" font-family="monospace" font-size="12">.</text>
      <rect x="200" y="15" width="50" height="40" rx="3" fill="#22c55e" stroke="#22c55e" stroke-width="2"/>
      <text x="225" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">kt+1</text>
    </g>
    
    <g transform="translate(0, 85)">
      <text x="175" y="0" text-anchor="middle" fill="#a855f7" font-family="monospace" font-size="10" font-weight="bold">Value Cache</text>
      <rect x="0" y="15" width="50" height="40" rx="3" fill="#a855f7"/>
      <text x="25" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v1</text>
      <rect x="55" y="15" width="50" height="40" rx="3" fill="#a855f7"/>
      <text x="80" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v2</text>
      <rect x="110" y="15" width="50" height="40" rx="3" fill="#a855f7"/>
      <text x="135" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v3</text>
      <text x="180" y="40" fill="#64748b" font-family="monospace" font-size="12">.</text>
      <rect x="200" y="15" width="50" height="40" rx="3" fill="#22c55e" stroke="#22c55e" stroke-width="2"/>
      <text x="225" y="40" text-anchor="middle" fill="white" font-family="monospace" font-size="8">vt+1</text>
    </g>
    
    <g transform="translate(350, 20)">
      <rect width="330" height="120" rx="8" fill="#f8fafc" stroke="#334155"/>
      <text x="165" y="20" text-anchor="middle" fill="#f59e0b" font-family="monospace" font-size="10" font-weight="bold">KV Cache Memory</text>
      <text x="10" y="50" fill="#f8fafc" font-family="monospace" font-size="9">Per layer: 2 x n_layers x n_heads x seq_len x d_head</text>
      <text x="10" y="75" fill="#60a5fa" font-family="monospace" font-size="9">Example: LLaMA-7B</text>
      <text x="10" y="95" fill="#94a3b8" font-family="monospace" font-size="8">32 layers x 32 heads x 2048 x 128</text>
      <text x="10" y="115" fill="#ec4899" font-family="monospace" font-size="8">Scales linearly with sequence length!</text>
    </g>
  </g>
</svg>

3. Multi-Query vs Grouped-Query Attention

<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="400" fill="#f8fafc"/>
  
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Multi-Head vs Multi-Query vs Grouped-Query Attention</text>
  
  <g transform="translate(30, 60)">
    <text x="130" y="0" text-anchor="middle" fill="#3b82f6" font-family="monospace" font-size="11" font-weight="bold">Multi-Head (MHA)</text>
    <rect x="10" y="15" width="50" height="25" rx="3" fill="#3b82f6"/>
    <text x="35" y="32" text-anchor="middle" fill="white" font-family="monospace" font-size="8">q1</text>
    <rect x="70" y="15" width="50" height="25" rx="3" fill="#3b82f6"/>
    <text x="95" y="32" text-anchor="middle" fill="white" font-family="monospace" font-size="8">q2</text>
    <rect x="10" y="50" width="50" height="25" rx="3" fill="#a855f7"/>
    <text x="35" y="67" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k1</text>
    <rect x="70" y="50" width="50" height="25" rx="3" fill="#a855f7"/>
    <text x="95" y="67" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k2</text>
    <rect x="10" y="85" width="50" height="25" rx="3" fill="#22c55e"/>
    <text x="35" y="102" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v1</text>
    <rect x="70" y="85" width="50" height="25" rx="3" fill="#22c55e"/>
    <text x="95" y="102" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v2</text>
    <text x="130" y="135" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="9">KV cache: 100%</text>
  </g>
  
  <g transform="translate(290, 60)">
    <text x="130" y="0" text-anchor="middle" fill="#a855f7" font-family="monospace" font-size="11" font-weight="bold">Multi-Query (MQA)</text>
    <rect x="10" y="15" width="50" height="25" rx="3" fill="#3b82f6"/>
    <text x="35" y="32" text-anchor="middle" fill="white" font-family="monospace" font-size="8">q1</text>
    <rect x="70" y="15" width="50" height="25" rx="3" fill="#3b82f6"/>
    <text x="95" y="32" text-anchor="middle" fill="white" font-family="monospace" font-size="8">q2</text>
    <rect x="40" y="55" width="120" height="25" rx="3" fill="#a855f7"/>
    <text x="100" y="87" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k (shared)</text>
    <rect x="40" y="90" width="120" height="25" rx="3" fill="#22c55e"/>
    <text x="100" y="122" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v (shared)</text>
    <text x="130" y="135" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">KV cache: 12.5%</text>
  </g>
  
  <g transform="translate(550, 60)">
    <text x="130" y="0" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="11" font-weight="bold">Grouped-Query (GQA)</text>
    <rect x="10" y="15" width="50" height="25" rx="3" fill="#3b82f6"/>
    <text x="35" y="32" text-anchor="middle" fill="white" font-family="monospace" font-size="8">q1</text>
    <rect x="70" y="15" width="50" height="25" rx="3" fill="#3b82f6"/>
    <text x="95" y="32" text-anchor="middle" fill="white" font-family="monospace" font-size="8">q2</text>
    <rect x="10" y="55" width="110" height="25" rx="3" fill="#a855f7"/>
    <text x="65" y="72" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k1 (group 1)</text>
    <rect x="130" y="55" width="110" height="25" rx="3" fill="#f59e0b"/>
    <text x="185" y="72" text-anchor="middle" fill="white" font-family="monospace" font-size="8">k2 (group 2)</text>
    <rect x="10" y="90" width="110" height="25" rx="3" fill="#22c55e"/>
    <text x="65" y="107" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v1 (group 1)</text>
    <rect x="130" y="90" width="110" height="25" rx="3" fill="#ec4899"/>
    <text x="185" y="107" text-anchor="middle" fill="white" font-family="monospace" font-size="8">v2 (group 2)</text>
    <text x="130" y="135" text-anchor="middle" fill="#f59e0b" font-family="monospace" font-size="9">KV cache: 25-50%</text>
  </g>
  
  <g transform="translate(50, 220)">
    <rect width="700" height="160" rx="8" fill="#f8fafc" stroke="#334155"/>
    <text x="350" y="25" text-anchor="middle" fill="#f59e0b" font-family="monospace" font-size="11" font-weight="bold">Comparison Summary</text>
    
    <text x="30" y="82" fill="#3b82f6" font-family="monospace" font-size="10">MHA: h key-value pairs per layer</text>
    <text x="30" y="80" fill="#a855f7" font-family="monospace" font-size="10">MQA: 1 key-value pair per layer (h x savings)</text>
    <text x="30" y="117" fill="#22c55e" font-family="monospace" font-size="10">GQA: g key-value pairs per layer (h/g x savings)</text>
    
    <text x="30" y="130" fill="#94a3b8" font-family="monospace" font-size="9">Quality: MHA > GQA > MQA</text>
    <text x="30" y="148" fill="#94a3b8" font-family="monospace" font-size="9">Memory: MQA > GQA > MHA</text>
    <text x="30" y="165" fill="#f59e0b" font-family="monospace" font-size="9">GQA: Best trade-off for modern LLMs</text>
  </g>
</svg>

4. Multi-Query Attention (MQA)

4.1 MQA Formulation

In MQA, all attention heads share a single key and value:

where each head uses the same :

4.2 Memory Savings

For attention heads:

• MHA: key-value pairs per layer
• MQA: 1 key-value pair per layer
• Savings: reduction in KV cache

4.3 Quality Impact

MQA can degrade quality, especially for:

• Complex reasoning tasks
• Long-range dependencies
• Tasks requiring diverse attention patterns

5. Grouped-Query Attention (GQA)

5.1 GQA Formulation

GQA divides heads into groups, each sharing a key-value pair:

where each group uses the same .

5.2 GQA Variants

5.3 GQA Training

Convert from MHA to GQA:

1. Mean pooling: Average K,V within each group
2. Fine-tuning: Train additional steps on large datasets

6. PagedAttention (vLLM)

6.1 The Memory Fragmentation Problem

Traditional KV cache allocation:

• Pre-allocate maximum sequence length
• Wastes memory for short sequences
• Causes OOM for long sequences

6.2 PagedAttention Concept

Inspired by operating system virtual memory:

• Divide KV cache into fixed-size blocks
• Allocate blocks on-demand
• No memory fragmentation

6.3 Block Management

Each block stores a fixed number of tokens:

where is the block size (typically 16 tokens).

6.4 PagedAttention Benefits

<svg viewBox="0 0 800 300" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="300" fill="#f8fafc"/>
  
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">PagedAttention: Virtual Memory for KV Cache</text>
  
  <g transform="translate(50, 50)">
    <text x="175" y="0" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">Physical Memory Blocks</text>
    
    <rect x="0" y="15" width="60" height="40" rx="3" fill="#22c55e"/>
    <text x="30" y="30" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Blk 0</text>
    <text x="30" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Seq A</text>
    
    <rect x="70" y="15" width="60" height="40" rx="3" fill="#22c55e"/>
    <text x="100" y="30" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Blk 1</text>
    <text x="100" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Seq A</text>
    
    <rect x="140" y="15" width="60" height="40" rx="3" fill="#3b82f6"/>
    <text x="170" y="30" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Blk 2</text>
    <text x="170" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Seq B</text>
    
    <rect x="210" y="15" width="60" height="40" rx="3" fill="#a855f7"/>
    <text x="240" y="30" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Blk 3</text>
    <text x="240" y="45" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Seq C</text>
    
    <rect x="0" y="65" width="60" height="40" rx="3" fill="#64748b"/>
    <text x="30" y="80" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Blk 4</text>
    <text x="30" y="95" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Free</text>
    
    <rect x="70" y="65" width="60" height="40" rx="3" fill="#64748b"/>
    <text x="100" y="80" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Blk 5</text>
    <text x="100" y="95" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Free</text>
  </g>
  
  <g transform="translate(400, 50)">
    <text x="175" y="0" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">Benefits</text>
    
    <text x="0" y="25" fill="#22c55e" font-family="monospace" font-size="10">No memory fragmentation</text>
    <text x="0" y="50" fill="#22c55e" font-family="monospace" font-size="10">Dynamic allocation (on-demand)</text>
    <text x="0" y="75" fill="#22c55e" font-family="monospace" font-size="10">Efficient memory sharing</text>
    <text x="0" y="100" fill="#22c55e" font-family="monospace" font-size="10">Supports variable-length sequences</text>
    <text x="0" y="125" fill="#22c55e" font-family="monospace" font-size="10">Copy-on-write for beam search</text>
    
    <text x="0" y="160" fill="#f59e0b" font-family="monospace" font-size="9">Block size: 16 tokens</text>
    <text x="0" y="180" fill="#f59e0b" font-family="monospace" font-size="9">Memory utilization: ~90%+</text>
    <text x="0" y="200" fill="#f59e0b" font-family="monospace" font-size="9">Speedup: 2-4x for batch inference</text>
  </g>
</svg>

7. Sliding Window Attention

限制每个 token 只关注固定窗口内的 token：

7.1 Complexity

7.2 Multi-Scale Sliding Window

Combine multiple window sizes:

8. Token Pruning

8.1 Attention-Based Pruning

Prune tokens with low attention scores:

8.2 Token Merging

Merge similar tokens instead of pruning:

9. Practical Optimization Strategies

9.1 KV Cache Quantization

Quantize cached keys and values:

Savings: 2-4x memory reduction with minimal quality loss.

9.2 Speculative Decoding

Generate multiple tokens in parallel:

1. Draft model generates tokens quickly
2. Target model verifies all tokens in parallel
3. Accept/reject based on probability

10. Summary

KV cache optimization is essential for efficient LLM inference. GQA provides the best balance of quality and efficiency, while PagedAttention enables high-throughput serving with near-zero memory waste.