【元來如此】第二章——打破序列長度限制，讓無限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.oeuau.cn/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 技術博客