Fast Transformer Decoding: One Write-Head is All You Need论文阅读笔记（MQA）

Motivation & Abs

增量推理对于MHA是非常慢的（难以并行），因为重复加载大的键/值会增大内存带宽的开销。为此作者提出了multi-query attention（MQA），其中不同注意力头共享相同的键和值，减小了增量解码的内存带宽要求。MQA可以大幅提升解码的速度，同时推理质量仅有略微下降。

Method

Multihead Attention (Incremental)

这里的计算次数是\(\Theta(bnd^2)\)，因为作者进行了如下简化假设：

\(m=n\)

\(k=v=\frac{d}{h}\)

\(n\leq d\)

带入可得使用einsum计算q的运算次数为\(bdhk=bdh\frac{d}{h}=bd^2\)，同理new_K和new_V也是同样的运算次数。计算logits的运算次数为\(bhmk=bhn\frac{d}{h}=bnd\leq bd^2\)，计算o的次数为\(bhmv=bhn\frac{d}{h}=bnd\leq bd^2\)，计算y的次数为\(bhvd=bh\frac{d}{h}d=bd^2\)。因此进行n次函数调用的运算次数为\(\Theta(bnd^2)\)。

进行n次函数调用的内存访问为\(\Theta(bn^2d+nd^2)\)。第一项来自K和V，第二项来自P矩阵。

通过以上观察可以发现，内存访问次数与算数运算次数的比率为\(\Theta(\frac{n}{d}+\frac{1}{b})\)，当\(n\approx d\)或者\(d\approx 1\)，瓶颈就变成了内存带宽。为了让增量推理更加高效，必须让这个比例远小于1。减小第二项只需要增加batch size，但渐小第一项比较困难。本文提出了一种方法，删除K和V的heads维度，同时在Q中保留这一维度。

Multi-Query Attention

MQA可以视作MHA的变体，不同的 head 共享一组键和值：

这样可以将KV cache的大小变为\(\frac{1}{h}\)，非常可观。