我们可以使用einops库有效地执行对张量的复杂操作。
首先,定义查询、键和值。...value张量的形状是一样的。...我们看看到底是怎么操作的
einsum帮我们做了两件事:
1、一个查询和键的矩阵乘法。...在我们的例子中,这些张量的形状是(1,4,2,256,64)和(1,2,256,64),所以沿着最后两个维度的矩阵乘法得到(1,4,2,256,256)。...GQA是最佳性能(MQA)和最佳模型质量(MHA)之间的一个很好的权衡。
下图显示,使用GQA,可以获得与MHA几乎相同的模型质量,同时将处理时间提高3倍,达到MQA的性能。