22倍加速還不夠，再來提升46%，而且方法直接開源！

這就是開源社區改進MIT爆火項目StreamingLLM的最新成果。

StreamingLLM可以在不犧牲生成效果、推理速度的前提下，實現多輪對話共400萬個token，22.2倍推理速度提升。

該項目在上線不到3個月時間內，GitHub項目標星達到5.7k star。

不過，StreamingLLM使用原生PyTorch實現，對于多輪對話推理場景落地應用的低成本、低延遲、高吞吐等需求仍有優化空間。

Colossal-AI團隊開源了SwiftInfer，基于TensorRT的StreamingLLM，可以進一步提升大模型推理性能46%，有效解決如上問題。

具體如何實現？一起來看。

開源地址：https://github.com/hpcaitech/SwiftInfer

StreamingLLM如何實現超長多輪對話？

大語言模型能夠記住的上下文長度，直接影響了ChatGPT等大模型應用與用戶互動的質量。

如何讓LLM在多輪對話場景下保持生成質量，對推理系統提出了更高的要求，因為LLM在預訓練期間只能在有限的注意力窗口的限制下進行訓練。

常見的KV Cache機制能夠節約模型計算的時間，但是在多輪對話的情景下，key和value的緩存會消耗大量的內存，無法在有限的顯存下無限擴展上下文。

同時，訓練好的模型在不做二次微調的前提下也無法很好地泛化到比訓練序列長度更長的文本，導致生成效果糟糕。

△圖來源：https://arxiv.org/pdf/2309.17453.pdf

StreamingLLM為了解決了這個問題，通過觀察了注意力模塊中Softmax的輸出，發現了attention sink的現象。

要知道，注意力機制會為每一個token分配一個注意力值，而文本最初的幾個token總是會分配到很多無用的注意力。

當使用基于滑動窗口的注意力機制時，一旦這幾個token被踢出了窗口，模型的生成效果就會迅速崩潰。

只要一直把這幾個token保留在窗口內，模型就能穩定地生成出高質量的文本。

比起密集注意力（Dense Attention）、窗口注意力（Window Attention）以及帶重計算的滑動窗口注意力(Sliding Window w/ Re-computing)，StreamingLLM基于attention sink的注意力機制無論是在計算復雜度還是生成效果上都表現優異。

在不需要重新訓練模型的前提下，StreamingLLM能夠直接兼容目前的主流大語言模型并改善推理性能。

SwiftInfer如何升級StreamingLLM？

Colossal-AI團隊的方法是，將StreamingLLM方法與TensorRT推理優化結合。

這樣的好處是，不僅能繼承原始StreamingLLM的所有優點，還具有更高的運行效率。

使用TensorRT-LLM的API，還能夠獲得接近于PyTorch API的模型編寫體驗。

具體來看，基于TensorRT-LLM可以重新實現KV Cache機制以及帶有位置偏移的注意力模塊。

如下圖所示，假設窗口大小為10個token，隨著生成的token增加（由黃色方塊表示），在KV緩存中將中間的token踢出，與此同時，始終保持著文本開始的幾個token（由藍色方塊表示）。

由于黃色方塊的位置會發生變化，在計算注意力時，也需要重新注入位置信息。

需要注意的是，StreamingLLM不會直接提高模型能訪問的上下文窗口，而是能夠在支持流式超多輪對話的同時保證模型的生成效果。

經過升級后，在原版本StreamingLLM基礎上，Colossal-AI團隊發布的SwiftInfer可帶來額外的最多46%的推理吞吐速度提升。

能為大模型多輪對話推理提供低成本、低延遲、高吞吐的最佳實踐。TensorRT-LLM團隊也在同期對StreamingLLM進行了類似支持。

團隊剛開源13B大模型

Colossal-AI目前已獲得GitHub星數三萬五千多顆，位列全球TOP400，細分賽道排名世界第一，可通過高效多維并行、異構內存等，降低AI大模型訓練/微調/推理的開發與應用成本，提升模型任務表現，降低GPU需求。作為主流開源AI大模型系統社區，Colossal-AI生態在多方面保持活躍更新。

值得一提的是，Colossal-AI最近還開源了另一個新模型Colossal-LLaMA-2-13B。

它僅用25B token 數據和萬元算力，效果遠超基 LLaMA-2 的其他中文漢化模型。

即使與其他采用中文語料，可能花費上千萬元成本，從頭預訓練的各大知名模型相比，Colossal-LLaMA-2在同規模下表現仍舊不錯。

13B 版本通過構建更為完善的數據體系，在知識性內容掌握程度，自然語言處理任務理解程度，以及安全性，價值觀等問題上，都有質的提升。

另外在云平臺方面，Colossal-AI云平臺在整合Colossal-AI系統優化和廉價算力的基礎上，近期發布了AI云主機的功能。

它能方便用戶以近似裸機的方式進行AI大模型的開發和調試，并提供了多種使用方式，包括：Jupyter Notebook、ssh、服務本地端口映射和grafana監控，全方位的為用戶提供便捷的開發體驗。

同時，還為用戶預制了含有ColossalAI代碼倉庫和運行環境的docker鏡像，用戶無需環境和資源配置，便可一鍵運行ColossalAI代碼倉庫中的代碼樣例。

Colossal-AI開源地址：https://github.com/hpcaitech/ColossalAI