大模型Transformer笔记:KV缓存

1 MHA(Multi-Head Attention)

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

  • 对于LLM来说,一般都是自回归地一个一个token的输出,也就相当于只有Transformer的decoder input在变化,之前作为prompt部分的k_{\le t}^{(s)},v_{\le t}^{(s)}是不变,可以缓存的(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

扫描二维码关注公众号,回复: 17402794 查看本文章

猜你喜欢

转载自blog.csdn.net/qq_40206371/article/details/143254221