不同repeat机制分析

简要介绍不同repeat在存储方面的机制差异，重点分析用通常方法实现的GQA为什么只能减小kv cache，但不能降低显存开销。

在numpy、pytorch等支持矩阵运算的库中都有repeat或类似的算子：

• numpy：np.repeat np.tile

• pytorch：torch.repeat_interleave torch.tensor.repeat

  注意torch的repeat和np的tile是对应的，与np中repeat对应的是interleave版本

• einops：repeat

以上numpy和pytorch的’repeat’都会实际地复制数据构造新矩阵，占用实际的内存空间。但einops的repeat经过了特定优化：在可能的情况下，只是重新构造一个张量头部返回（即类似.view返回的结果，此时的新头部和原头部共享相同的数据存储区域），这种情况下就能够节省空间。参考https://github.com/arogozhnikov/einops/issues/183 和 https://github.com/arogozhnikov/einops/issues/202

“可能的情况”其实就是能够通过view调整（调整头信息中的stride即可）的情况，在通常方式实现的GQA中，代码可能类似这样（g就是GQA的组数）：

from einops import rearrange, repeat
...
q = rearrange(q, "b (h g) t d -> b h g t d", g=self.g)
k = repeat(k, "b h t d -> b h g t d", g=self.g)
v = repeat(v, "b h t d -> b h g t d", g=self.g)
...

这种变化会导致内存不连续，因此einops会自动复制+构建新数据。因此这部分的显存开销和无GQA时是一样的。