【元来如此】第二章——打破序列长度限制，让无限Token成为可能！
正文共：1929字 6图
预计阅读时间：3分钟
作者：思成
软件生态中心·应用平台部
“
此前，我们在【元来如此】第一章——大模型技术 · 起航&推理篇中对大模型推理技术做了一些基础介绍。
还没有关注的同学快来补课!
(点击图片，即可跳转阅读)
目前主流大模型推理框架比如vLLM、FasterTransformer+Triton Server、TGI（Text Generation Inference）等在Latency（时延）、Throughput（吞吐）、Utilization（利用率）、易用性等方面各有优劣。但是当前主流大模型只能生成有限长度的文本（比如LLaMA 2 7B模型的限制是2048个Tokens），剩余的部分会被截断掉。
本文，针对以上长序列问题，我们将介绍Tecorigin在无限长序列推理生成上的一些探索和努力。
我们基于StreamingLLM[1]算法提出了inference-infinitely，并在Teco-Inference-Engine（推理引擎框架）上支持了这种算法，从而赋予了大模型无限长度生成的能力。我们在Baichuan2-7B-Chat上进行了验证。
StreamingLLM
图1 StreamingLLM和其他方法示意图，引用自[1]
1(a)
如图1(a)所示，是Transformer模型最基础的Attention形式，这种Dense的方式在生成第T个Token的时候需要关注前面0到T-1个Token，做全局的注意力，这意味着0(T2)的时间复杂度和不断增加的KV cache显存占用。
1(b)
1(b) 其改进版本，一种Window Attention，这种方式在生成第T个Token的时候只需要关注前面L个Token即可，这种方式获得了0(TL)的时间复杂度，以及可控的KV cache显存占用，但是也带来了断崖式的效果下降。
1(c)
1(c) 另一种改进版本，称作Sliding Window Attention，需要在每次计算的时候重新计算长度为L的KV cache，这种方法会有比较好的模型效果，但是也带来了性能的损失。
1(d)
1(d) StreamingLLM的做法，理解起来很简单，在1(b)的基础上，保留一些初始的Token。通过这样的方式即有0(TL)的时间复杂度，同时也取得了比较好的模型效果。
StreamingLLM的方法到底为什么生效呢？其实之前在Lost in the Middle [2] 中的研究已经发现，Dense Attention的方式其注意力会集中在序列中的头和尾两个部分，对于中间Token的Attention会比较弱。
图2 lost in the middle引用自[2]
这种现象在StreamingLLM中也得到了进一步的验证：
图3 除去0和1层的Head之外，其他层都重点关注头部部分，引用自[1]
论文中也对此进行了一定解释：要么是开始的Token在模型推理过程中扮演着重要的语义信息，要么是模型出于某种原因，导致了对于这些位置的归纳偏置，即bias。之后论文通过实验进一步验证了后者更可能是主要原因。
基于上面的现象，StreamingLLM中保留了一定的初始Token，通过这种Attention sink的方案稳定Attention计算。之后通过Rolling KV的方式保留一定数量的最新Token。
如图4所示。
图4 Streaming LLM中的KV cache方案，
引用自[1]
到此为止，我们从介绍了StreamingLLM的来龙去脉，下面让我们深入细节，进一步思考。
KV cache
上一节我们多次提到了KV cache这样的字眼，简单说：KV cache是一种在自回归大模型下的优化方法，该方法通过缓存一定的过去信息（K和V），用于加速当前step模型的推理过程，避免重复计算。
详细来看，如图5所示：
Step 1：正常的Attention计算，之后将K矩阵和V矩阵的结果缓存下来KcacheVcache。
Step 2：根据第一步生成的Token，生成当前step的Qstep/Kstep/Vstep向量（图5中Step N蓝色部分），并将Kstep/Vstep和第一步生成的KcacheVcache拼接到一起生成真正的K/V向量，从而进行后面的计算。
图5 KV cache机制，引用自[3]
根据这样的原理，不难得到KV cache在显存中的占用：2（K+V）*L（模型层数）*B（batch大小）*S（序列长度）*D（隐层纬度）*2（fp16字节），以LLaMA 2 7B模型为例带入得到2*32*1*4096*4096*2=2GB。
让我们回想一下刚刚介绍的StreamingLLM算法，当超过最大长度之后，需要Rolling KV从而生成下一个Token，这样势必会造成大量的KV cache读写操作，给本就访存密集的推理过程进一步增加大量额外访存量，大大降低推理性能。
带着这样的疑问，我们深入StreamingLLM作者给出的开源实现方案。实际上，从代码中可以看到是如下流程：
Step 1：通过设置一个max_gen_len参数控制了本次推理最大的生成序列长度，并和当前prompt长度一起重新开辟了KV cache最大的显存占用空间。
Step 2：对KV cache空间进行一次整体的Rolling。
Step 3：后续正常推理过程。
可以看到这样的实现方法解决了频繁的KV cache Rolling带来的访存问题，但是在每次问答中并不能进行无限长度输出。
Inference-infinitely
基于上述的问题，我们提出了Inference-infinitely，并在Teco-Inference-Engine上进行了实现，提供了每次问答都可以无限输出的能力。
图6 Inference-infinitely示意图
图6中描述了Inference-infinitely的算法：
Step 1：保留初始Token（黄色）。
Step 2：当达到本次输出的最大长度之后，将KV cache通过参数W（图中为128）进行一次整体滑动（蓝色），并生成新的Token（红色）。
Step 3：每滑动一次提供W个Token生成的空间（灰色），当buffer填满之后，再进行Step 2。
通过这样的算法，很好的解决了上述提到的KV cache Rolling带来的频繁访存，以及每次问答的生成都可以有无限的输出长度，同时也保证了模型的推理效果。
至此，本文介绍了Tecorigin 在大模型推理——无限长序列生成上的一些探索和努力。未来，会有更多的大模型技术跟大家一起分享、交流、讨论。
免费试用申请
如果您对我们的产品感兴趣，
可点“http://www.tecorigin.com/cn/developer.html ”，进行试用申请。
参考文献
[1] [2309.17453] Efficient Streaming Language Models with Attention Sinks (arxiv.org)
[2] [2307.03172] Lost in the Middle: How Language Models Use Long Contexts (arxiv.org)
[3] 使用 FasterTransformer 和 Triton 推理服务器加速大型 Transformer 模型的推理 - NVIDIA 技术博客