LLM 20,000 QPS背后的力量：揭秘Character AI的高效推理技术_character.ai cross layer attention-CSDN博客

本文链接：https://blog.csdn.net/xx_nm98/article/details/142576760

https://athekunal.medium.com/character-ai-optimized-inference-blog-post-explained-ce192761536d

摘要

大型语言模型（LLM）推理吞吐量的主要瓶颈在于注意力键值（KV）缓存的大小。它不仅决定了GPU上能够容纳的最大批次大小，还主导了注意力层的I/O成本。我们采用了以下技术，成功将KV缓存大小减少了20倍以上，同时保证了质量不受损。

LONGFORMER
Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
RadixAttention from Structured Generation Language paper

在深入探讨本文之前，我建议您观看视频,以便更直观地了解KV缓存：

https://www.youtube.com/watch?v=Mn_9W1nCFLo&t=2422s
https://www.youtube.com/watch?v=80bIUggRJf4

如果您希望更深入地理解Transformer模型，我还推荐您观看视频：

https://www.youtube.com/watch?v=kCc8FmEb1nY&t=3868s&pp=ygUSZ3B0IDIgZnJvbSBzY3JhdGNo
https://www.youtube.com/watch?v=rPFkX5fJdRY&pp=ygUbY29kZSBlbXBvcnVpdW0gdHJhbnNmb3JtZXJz
https://www.youtube.com/playlist?list=PLQWPycXvOsB-2rylhr9ltoWuvybuGKBHD

LONGFORMER

在Transformer层中，注意力计算与序列长度呈二次方关系增长(编码器)。在仅解码器的Transformer模型中，为了生成下一个标记，我们需要计算与所有先前标记的注意力分数。这解释了为什么大型语言模型（LLMs）中存在上下文限制，例如GPT-4的上下文限制约为128K，而像Mistral这样的开源模型则约为32K。

然而，更复杂的任务需要超出这些限制的上下文大小。LONGFORMER论文通过滑动窗口注意力机制解决了日益增长的上下文大小问题。

全面注意力机制（a）：这种注意力机制是全连接的，每一个位置（token）都会与序列中的所有其他位置进行计算。图中绿色的方格表示每个位置都参与了注意力计算。这种方法虽然能捕捉到全局信息，但计算和存储的开销是O(n^2)，其中 n 是序列的长度。

滑动窗口注意力机制（b）：在序列中为每个位置只计算一个固定窗口内的注意力（如图中的对角线带状区域）。这样可以降低计算复杂度，但只能捕捉局部信息，适用于局部依赖关系强的任务。

扩张滑动窗口注意力机制（c）：扩展滑动窗口注意力机制引入了一个稀疏化（dilation）的概念，通过增加窗口的跨度来覆盖更广的范围（但不是连续的），如图中的网格状图案。这种方法在减少计算量的同时，保留了一定程度的全局信息捕捉能力。

全局滑动窗口注意力机制（d）相较于传统的滑动窗口注意力机制更进一步，它根据任务需求考虑了先前的关键标记。若你面对的是问答任务，可将问题标记纳入全局注意力范畴，鉴于问题通常简短，其注意力计算将接近于O(~n)复杂度，确保模型在生成过程中始终聚焦于问题或指令。对于分类任务，则可将[CLS]标记设为全局注意力，使模型明确需对句子进行分类。对于BERT模型，若涉及掩码语言模型预训练任务，可将[MLM]标记作为生成任务的全局标记。

若你使用的是仅解码器的语言模型，那么可以采用逐步扩大的滑动窗口注意力机制。例如，在Transformer的第一层中，模型可能仅关注其前10个标记，第二层则扩展到20个，依此类推。这一设计灵感来源于卷积神经网络中的感受野概念。此外，你还可以采用扩张滑动窗口注意力来扩大有效感受野。

在这里插入图片描述

Transformer的注意力计算具有平移不变性，因此它不依赖于标记的位置。在原始论文中，作者向标记中添加了位置编码。在《LONGFORMER》论文中，作者在全球+滑动窗口注意力机制中再次复制了位置嵌入，且上下文长度与Transformer相匹配。例如，作者还通过初始化和复制ROBERTA模型的位置编码权重，并在滑动窗口注意力中重复这些权重，来训练他们自己的LONGFORMER模型（LED）。

我还想提及另一篇关于注意力汇聚点的论文<>。在这篇论文中，作者在一个仅解码器的模型中实现了滑动窗口注意力，但仅将前四个标记作为全局注意力。然而，前四个标记并不提供任何上下文信息，它们仅仅是位置标记。当作者用“\n\n\n\n”替换前四个标记时，他们得到了相似的结果，这表明选择前四个标记仅是一个位置要求，而非语义要求。

跨层注意力

从左至右，首先介绍的是多查询注意力（Multi-Query Attention，MQA）。在这里，方框代表权重投影层，它们负责将输入转换为查询（Query, Q）、键（Key, K）和值（Value, V）。接下来，我们将遵循多头注意力（Multi-Head Attention, MHA）的工作原理进行阐述。

符号说明：d：输入层的维度；nH：头的数量（本例中假设为4）。

多头查询注意力

对于序列长度为1的情况，我们将标记嵌入转换为查询（Q）。对于键（K）和值（V），我们进行类似的转换，然后计算多头注意力（MHA）的自注意力分数。然而，这种方法会消耗大量内存，总参数大小为：3 (For Q, K and V) * d * d/nH * nH = 3d²。

2017年，最初的Transformer论文问世，当时使用的训练环境是P100 GPU，其双精度浮点运算能力为5.3 TFLOPs，内存为16GB。而相比之下，如今的A100 GPU拥有80GB的GPU内存和9.7 TFLOPs的fp64计算能力。GPU的计算能力显著提升，但内存的增长却相对温和。因此，现代GPU面临着内存瓶颈的问题。由于GPU在处理数据时，如果数据更接近其内部（根据GPU内存层次结构），则计算速度更快，所以内存限制成为了一个挑战。

此外，在自回归模型中，我们有一个KV缓存（即先前生成的标记的键和值矩阵），当前生成的标记会关注所有这些先前的KV缓存以预测下一个标记。与其计算K和V矩阵，我们不如将它们缓存起来。但这种方法内存消耗大，因为缓存的大小会持续增长，加上内存瓶颈的限制，我们受到其不断增长的大小的制约。

因此，引入了多查询注意力（MQA）和分组查询注意力（GQA）来减少内存开销。在MQA中，我们将所有查询头组合成一个组，并使用一个键和值矩阵（而不是多头注意力中的nH个键和值矩阵）。同样地，为了将nH个键头合并为一个，我们对所有键权重矩阵进行逐元素平均处理。

多查询注意力机制

在上面的图示中，您可以看到所有头的权重矩阵都进行了平均并合并为一个。最终，我们不再有四个不同的输出进行拼接，而是有一个相同的输出进行拼接。值矩阵的处理方式也是如此。不过，现在我们只需存储一个权重矩阵，而不是四个。因此，总参数大小为：

查询（Query）的参数大小为 d * d/nH * nH；
键（Key）和值（Value）的参数大小各为 d * d/nH，两者相加为 2 * d * d/nH；

所以，总参数大小为 d² + 2d²/nH。

分组查询注意力机制

对于分组查询注意力机制，我们不是将所有查询合并为一个组，而是将它们分成多个组，并在每个组内部进行聚合。

如果我们将KV分为两组，那么我们就一次将两个键分为一组并重复它们。从经验上看，这种方法表现更佳，并且被所有最先进的开源模型（如Llama、Mistral和Qwen）所采用。与多头注意力（MHA）相比，这种方法减少了KV缓存的内存占用，因为我们需要存储的矩阵数量更少了。

跨层注意力机制更进一步，通过跨Transformer层绑定权重来实现，在跨层注意力机制中，你还将在各层之间共享KV矩阵。

这里我们可以选择将层分为2组（CLA2）或3组（CLA3）。需要注意的是，在将层分为3组时，会有一部分层无法完整地归入某一组。作者尝试了不同的变体，将剩余的层单独归为一个组，置于顶部或底部。此外，作者还从头开始训练了他们的模型，并报告了验证困惑度。

让我们以H128-GQA4为例其命名规则是H<嵌入维度>-<模型><组数>-CLA<层组数>。

因此，H128-GQA4拥有20个Transformer层，所以键值层（KV层）的数量也是20。为了计算所需的字节数，我们可以这样计算：2（键和值）* 128（嵌入维度）* 4（键值头数）* 20（键值层或层数）* 2（float16占用16位，即2字节）= 40,960字节。

在Transformer层中，头的数量通常为16（因为给定的Query头数量为16）。所以GQA4表示我们有4组Key头，因此键值头的总数 = 16/4，即4。GQA2则表示我们有两组Key头，每组包含8个权重矩阵。而GQA1与MQA相同，意味着所有Key头都归为一组。

接下来，我们来计算CLA模型的参数大小。

根据我们原始的命名规则H256-GQA4-CLA2，我们可以这样计算参数大小：2（键和值）* 256（嵌入维度）* 4（键值头数）* 10（Transformer层数，实际为20层，但因为是CLA2配置，所以按10层计算）* 2（float16占用16位，即2字节）= 40,960字节

现在，如果我们有H128-MQA-CLA2的配置，计算方式如下：2（键和值）* 128（嵌入维度）* 1（键值头数，所有键/值头归为一组）* 10（Transformer层数，实际为20层，但因为是CLA2配置，所以按10层计算）* 2（float16占用16位，即2字节）= 10,240字节

在KeepEnds架构中，较高的Transformer块不共享权重。因此，如果我们有20层Transformer，那么其中的18层（即第2层到第19层）将共享9个KV矩阵，而剩下的两层（第1层和第20层）则各自拥有1个（即1+1）。这样总计为1+9+1=11个（如图所示）。

在DenseFront架构中，20层Transformer中的前10层仅共享一个KV投影，而后10层则各自拥有独立的KV投影。因此，总计为1（前10层共享1个KV投影）+ 10（后10层各自的KV投影）。

在DenseBack架构中，情况则相反，20层Transformer中的后10层仅共享一个KV投影，而前10层则各自拥有独立的KV投影。因此，总计为10（前10层Transformer各自的KV投影）+ 1（后10层共享1个KV投影）。

作者发现，在帕累托前沿的内存与准确度的权衡下，H128-MQA-CLA2配置表现更佳。当他们将模型从10亿参数扩展到30亿参数时，发现这种配置的性能有了显著提升。这种权重共享极大地减少了Transformer模型的内存占用。

RadixAttention

第三篇论文主要基于结构化语言生成（SGLang）。其中，我们将讨论论文中关于基于基数树的键值（KV）缓存的部分。

如前所述，Character.ai是一家角色扮演公司，能够以特定角色的风格回答问题。为了保持角色的连贯性，聊天历史需要包含所有系统提示（即角色的基本定义和角色设定）以及聊天内容。

键值缓存对于每个句子或前缀都是独特的，因为它基于上下文构建。因此，如果我们能够缓存聊天历史和系统提示，就可以实时解码令牌（tokens）。

关键在于，我们需要将键值缓存张量作为值存储，而键则是句子或前缀。但大多数聊天内容和系统提示都共享相同的系统提示和角色设定。或许当你切换到不同角色时，角色描述会发生变化，但系统提示保持不变。这构成了一个树状层次结构。

在这里，我们不能简单地采用最近最少使用（LRU）缓存策略，因为大多数提示共享相同的前缀、系统提示或聊天历史。因此，我们可以构建一个TRIE数据结构，其中边代表当前对话，如果对话具有相同的前缀，则它们共享一个节点。

上述图表解释得非常到位。您可以看到，对话是分层级的，如果某个叶节点不活跃，那么它就会被从缓存中移除（基于某些启发式规则）。这有点像您从缓存记忆中删除内容，如果您重新开始了这些被移除的对话中的任何一个，您可以重新计算键值缓存。

给定一个提示和系统消息，您可以遍历这个基数树，并依次将您的提示与基数树中的边进行匹配（优先匹配最长的共享前缀），最后到达一个节点。在那里，您需要关注您的图遍历路径，并生成下一个令牌。

随着Character.ai为用户提供服务并个性化内容，它可以为每位用户构建这样的树结构并存储键值缓存。此外，为了连续批处理，最好将具有相同前缀和系统提示的查询进行分组，这样您就不必为每个查询反复遍历基数树了。构建基数树只会产生0.2秒的额外开销，考虑到其带来的好处，这是可以接受的。

建议在提示的开头加上类似“你是一名财务分析师……”的系统提示，这样我们就可以根据系统提示将它们分组，然后遍历树结构。然而，作者采用了一种不太高效的方法来重新格式化提示，即让GPT-4重新排列提示，这会产生一定的延迟。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述