LLM 推理加速：decode 阶段的 Attention 在 GPU 上的优化

作者：董纪莹

随着大语言模型（Large Language Models，LLMs）在各领域的广泛应用，如何以低成本构建高吞吐、低延迟的推理服务成为了一个紧迫的问题。考虑到 LLM 在 GPU 上推理时参数量和计算量较大以致于单流执行就可以充分利用 GPU 资源，我们可以把 LLM 的推理延时分解到 kernel level，因此，进一步的，不考虑时间占比小的 kernel 计算后，LLM 的延时优化也就相应的分解成 GEMM 和 Attention 的 kernel 优化。

RTP-LLM 是阿里巴巴智能引擎团队开发的大模型推理加速引擎，作为一个高性能的大模型推理解决方案，它已被广泛应用于阿里内部。在这篇文章里，我们将基于 RTP-LLM 的实践，介绍 decode 阶段的 Attention 在 GPU 上是如何优化的。

一、背景

我们比较熟悉的 Attention 计算如下图所示，包含 Q 与 K 相乘，其结果在 mask 后做 softmax，然后与 V 相乘，得到 Attention 的结果。在 LLM 推理的 decode 阶段，由于加入 KV Cache 优化，一次迭代只需要计算新增的一个 token，因此计算可以变化为当前 step 的 Q（seq == 1）与 K Cache、V Cache 做计算。







计算过程中各 tensor 的 shape 可以表示为：、

参数的解释如下表：

在本文的分析中，我们考虑简单的 Multi Head Attention 实现，即 H == H_kv。