大模型Transformer笔记：KV缓存

1 MHA（Multi-Head Attention）

• 最经典的多头注意力
  •  等价于多个独立的单头注意力的拼接
  • 

• 对于LLM来说，一般都是自回归地一个一个token的输出，也就相当于只有Transformer的decoder input在变化，之前作为prompt部分的是不变，可以缓存的（KV cache)
  • 而KV cache的减少可以让我们有更长的context prompt，更快的推理速度，更低的推理成本

2  MQA(Multi-Query Attention)

Fast Transformer Decoding: One Write-Head is All You Need 2019

PaLM [6]、StarCoder [7]、Gemini [8]

• 所有注意力头共享同一套K,V
  • 
  • ——>KV 缓存减少到1/h
  • KV参数的减少可以到FFN/GLU规模的增大来弥补

3 GQA（Grouped-Query Attention）

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints emnlp 2023

所有 Head 分为 g个组（ g可以整除 head数量 h）

LLAMA2-70B ， LLAMA3 

参考内容：缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA