什么AI應用每秒處理20000個AI推理請求，達到2024年谷歌搜索流量的1/5?

　　答案是獨角獸Character.ai，由Transformer作者Noam Shazeer(后面簡稱沙哥)創辦。

　　剛剛，沙哥公布了推理優化獨門秘訣，迅速引起業界熱議。

　　具體來說Character.ai在整個服務堆棧中實現了如下成績:

　　內存高效架構設計:將KV緩存大小減少20倍以上，而不會降低質量

　　Attention狀態緩存:95%請求無需重算

　　直接用in8精度量化訓練:推理零損失還省顯存

　　Character.AI通過以上種種優化，已經把推理成本降低到最初的1/33，如果用市場上最好的商業API來支撐這種級別的流量，成本會比現在高出13.5倍!

　　眾多公布的方法中，原生int8訓練是最受關注的。

　　雖然大多數技巧都來自公開研究，但是正如網友所說，知道如何把它們高效整合在一起實現的團隊才是真正的護城河。

　　秘訣1:高效利用顯存，attention 參數量降低20倍

　　大模型的一大痛點是顯存占用高，導致無法支持大批量推理。Attention 層中的 Key-Value(KV)緩存便是罪魁禍首之一。

　　為了降低顯存占用，Character.AI在Attention層大動手術:

　　全面采用MQA(Multi-Query Attention)

　　與大多數開源模型中采用的GQA(Grouped-Query Attention)相比，將KV緩存大小減少了8倍。

　　而MQA正是沙哥本人2019年在谷歌期間提出的，有網友評價“當一個人能在生產環境中引用自己的論文，就達到了一個新的高度”。

　　混合注意力視野

　　將局部注意力與全局注意力層交織在一起，使用滑動窗口訓練局部注意力，將復雜度從 O(length^2) 降低到 O(length)。

　　團隊發現，將大多數注意力層的注意力范圍減少到1024不會對評估指標產生重大影響，包括長上下文大海撈針基準。在Character.ai生產模型中，每6層中只有1層使用全局注意力。

　　跨層KV共享

　　團隊將KV緩存綁定在相鄰的注意力層上，這進一步將 KV緩存大小減少了2-3倍。

　　對于全局注意力，跨塊綁定多個全局層的KV緩存，因為全局注意力層在長上下文用例中主導KV緩存大小，團隊發現跨層共享KV不會降低質量。

　　下圖中左半部分是標準Transformer設計，每個注意力都是全局注意力。右半部分為Character.ai的設計，藍色框表示全局注意力，綠色框表示局部注意力，連線表示KV共享。

　　這一套組合拳下來，KV緩存大小減少20倍以上，顯存再也不是瓶頸了。

　　秘訣2:巧用狀態緩存，95%請求無需重算

　　Character.AI還有一招神來之筆，就是在不同對話之間緩存Attention狀態。

　　作為聊天機器人角色扮演服務，Character.AI上大部分對話都是連續多輪的，平均每個對話包含180條消息。如果每次都要重新計算前面的狀態，成本可想而知。

　　于是團隊設計了一個緩存機制，把每個對話的Prefix和生成的消息都緩存在內存中，供后續調用。

　　借鑒RadixAttention的思路，樹狀結構的LRU緩存組織緩存的KV張量。緩存的KV值由前綴token的Rolling Hash速檢索最長匹配的緩存，即使前綴只有部分匹配也能命中。

　　更妙的是，他們還用會話保持(Sticky Session)把同一對話路由到同一個服務器，進一步提高緩存命中率。最終做到95%的請求都能復用已有緩存，大幅降低了計算成本。

　　下圖中，藍色框表示主機內存上的緩存張量。綠色和黃色框表示CUDA內存上的KV緩存。當新查詢到達時，它檢索最長匹配前綴的KV緩存，Rolling Hash系統允許檢索部分匹配消息的緩存。

　　秘訣3:直接量化訓練，推理零損失還省顯存

　　最后一招，Character.AI沒有采用常見的“訓練后量化”，而是直接用Int8精度訓練模型。

　　這種格式雖然表達精度降低，但通過精心設計定制的矩陣乘和 Attention 內核，不僅把訓練效率提高了好幾倍，而且還能無損用于推理。

　　不過沙哥在這里暫時留了一手，表示“量化訓練本身就是一個復雜的話題，將在以后的文章中繼續討論。”

　　沙哥其人

　　最后再來介紹一下傳奇人物Noam Shazeer本人。

　　他1994年拿了信息學奧賽IOI金牌，后來畢業于杜克大學。

　　2000年加入谷歌，當時全公司只有200人左右，他參與了谷歌搜索的拼寫糾正功能，后來也負責過早期廣告系統。

　　據知情人透露，在當初面試谷歌時，沙哥就被問到如何實現拼寫糾正。他描述了一種根據其他用戶的輸入輸入記錄，進行統計驗證的方法。

　　面試官Gmail之父Paul Buchheit意識到，沙哥的方案比谷歌當時使用的要好。沙哥成功入職之后就把他的面試方案寫出來了。

　　在Transformer開山之作《Attention is All You Need》研究中，沙哥最后一個加入團隊，一來就負責重新編寫了整個代碼。

　　在沙哥出手之前，Transformer早期原型性能并沒有超越當時流行的LSTM方案，是他把早期設計中的卷積等模塊都拿掉，給出了一個極簡主義方案。最終破了BLEU測試的記錄，同時計算效率也更高。

　　隊友用“他是一個巫師”來評價他的工程和代碼能力。

　　除此之外，沙哥還有驚人的遠見。在Transformer架構問世不久，他就給谷歌高層寫信，提議公司放棄整個搜索索引，并用Transformer架構訓練一個巨大的神經網絡替代。

　　2021年，沙哥離開谷歌后創辦了Character.AI，讓玩家簡單自創個性化AI陪聊，目前估值約50億美元。

　　最近有消息稱，Meta與馬斯克的𝕏都在爭取與他們合作，把聊天機器人引入社交平臺。