近期大模型領(lǐng)域里最火的熱詞,或者說技術(shù)創(chuàng)新點,非Engram (DeepSeek最新論文里設(shè)計的Engram機制) 莫屬。今天我們想分享的,是英特爾圍繞Engram開展的早期探索——用至強® 處理器獨立運行整個Engram模塊,并使用其內(nèi)置的英特爾® AMX(高級矩陣擴展)技術(shù)對其進行加速的初步成果或收獲。

我們希望這次分享,或能作為參考,或是作為開端,能為未來Engram以及集成它的大模型的部署和實踐,拓展和探明更多可能性及隨之而來的潛在應用優(yōu)勢。

Engram設(shè)計初衷:

讓大模型走向“查算分離”

讓我們先簡單回顧Engram的源起,它出現(xiàn)在公眾視野,是源自DeepSeek聯(lián)合北京大學發(fā)布的論文《Conditional Memory via Scalable Lookup》。業(yè)界對它的評價,是為“破解萬物皆推理”模式引發(fā)的大模型的記憶困境提供了全新思路。這里提到的記憶困境,指的是寶貴的算力被消耗在本可直接調(diào)取的靜態(tài)知識檢索上,這不僅會拖慢響應速度、增加推理成本,還讓大模型在復雜任務上的性能突破陷入瓶頸。

該論文創(chuàng)新地在大模型中提出了“查算分離”理念,通過Engram引入外置記憶模塊,將“靜態(tài)、常見、局部”的知識從計算里解放出來,在拉升檢索效率至O(1)復雜度的同時,也把寶貴的算力資源留給Transformer專心做上下文理解和推理。

Engram核心創(chuàng)新:

用“外置記憶模塊”實現(xiàn)查算分離理念

Engram將“查算分離”落到實處的做法,就是把大模型里的“計算”和“超大規(guī)模記憶”解耦,Transformer的算子全部在GPU/加速卡上計算,而龐大的Engram Embedding表放在CPU內(nèi)存或高速存儲設(shè)備上存查。如圖1所示,GPU與CPU分工合作并通過異步方式協(xié)同,GPU執(zhí)行前一步計算的同時,CPU可提前預取后續(xù)計算所需的N-gram Embedding表,當計算執(zhí)行到“Transformer Block with Engram”時,所需的靜態(tài)知識已經(jīng)就位。

這種分工模式改變了傳統(tǒng)大模型“推理既要計算又要記憶”的狀態(tài),就像給學者配備了一本可即時查閱的百科詞典,無需每次都從頭推導基礎(chǔ)知識點,而是將精力集中在深度思考上。

圖1 大模型通過集成Engram實現(xiàn)查算分離

同時基于這一設(shè)計,Engram能充分利用CPU平臺對大容量內(nèi)存的有效支持,將“超大規(guī)模記憶”部分卸載至CPU平臺上。論文數(shù)據(jù)提到:“將 1000 億參數(shù)的表卸載至CPU內(nèi)存時,僅產(chǎn)生可忽略的開銷(小于 3%)。” 這表明,Engram 能有效突破GPU顯存限制,為大規(guī)模參數(shù)擴展提供支持,從而為大模型用戶帶來顯性收益。i

Engram計算流程解析

如圖2所示,在Engram的計算架構(gòu)中,其被嵌入Transformer主干網(wǎng)絡(luò),有以下幾個主要工作階段:

圖2 Engram計算架構(gòu)及核心工作流程

  • N-gram哈希檢索階段:模塊對輸入序列提取2-gram、3-gram等不同子序列(如圖中的“the Great”、“Alexander the Great”),每個子序列對應其階數(shù)的多頭哈希機制,從預構(gòu)建的靜態(tài)N-gram 嵌入表(如圖中的2-Gram Embedding、3-Gram Embedding表)中以O(shè) (1)復雜度查詢對應嵌入。
  • 動態(tài)門控融合階段:N-gram 嵌入表先經(jīng)Concat(通道拼接)形成融合記憶向量,再通過兩個Linear(線性轉(zhuǎn)換)層分別投影為Key(記憶語義摘要)與Value(待注入信息)。隨后將當前Transformer層的Input Hidden(全局上下文)與 Key 做 Scaled Dot Product (點積計算)生成門控權(quán)重,動態(tài)過濾與上下文無關(guān)的Value信息。加權(quán)后的Value再通過Conv卷積計算完成局部融合。
  • 殘差集成階段:局部融合后的記憶特征通過加法操作,殘差加回該Transformer Block的輸入,并直接輸入后續(xù)的Attention與MoE層。

英特爾的探索:

用CPU獨立運行及加速整個Engram模塊

在Engram相關(guān)論文發(fā)表,DeepSeek開源上述流程的Demo代碼后,不少業(yè)內(nèi)專家和機構(gòu)都開展了相關(guān)的復現(xiàn)、驗證及測試工作。我們的探索則更進一步——不同于原論文中Concat之后的工作任務將交還給GPU執(zhí)行,我們不僅將N-gram哈希檢索階段的計算放在至強® 平臺上執(zhí)行,還將動態(tài)門控融合階段中的Linear轉(zhuǎn)換計算和Conv卷積計算也放到該平臺上運行。換言之,我們是基于充分的性能調(diào)優(yōu),嘗試用CPU平臺獨立運行和加速整個Engram模塊。

這種“更進一步”的底氣,來自AMX技術(shù),這是從第四代英特爾® 至強® 可擴展處理器開始就內(nèi)置于至強® 處理器,且到目前為止也是全球主流服務器CPU產(chǎn)品中僅為該產(chǎn)品線所獨有的CPU內(nèi)置型矩陣計算加速技術(shù)。無論是Linear轉(zhuǎn)換計算還是Conv卷積計算,都屬于矩陣密集型計算,因此,理論上內(nèi)置AMX的至強® 處理器可以“順勢”完成整個Engram的運行和加速。

圖3 從Engram原有工作流程轉(zhuǎn)向用CPU運行整個模塊并導入AMX加速

基于此,我們著手在Linear轉(zhuǎn)換計算和Conv卷積計算中利用AMX技術(shù)開展了性能加速。具體來說,在Engram中,兩個Linear轉(zhuǎn)換計算分別需要處理大量維度映射的矩陣乘法,而AMX的專用矩陣計算單元可并行處理多批次、長序列的高維矩陣乘法,且單條指令可完成更大規(guī)模的矩陣運算,運算效率遠超傳統(tǒng)CPU計算或向量計算(如AVX-512)。在Conv的 short_conv(短卷積)計算中,AMX的矩陣運算能力也可針對短卷積的“小窗口、高并行” 特性實現(xiàn)優(yōu)化。此外,AMX還對BF16/FP16/INT8等不同的數(shù)據(jù)格式有著良好支持,能進一步提升矩陣運算加速的性能與靈活性。

我們目前已完成了一些初步測試,如圖4和圖5所示,其結(jié)果表明,在同一款至強® 6處理器平臺上,如果使用AVX-512加速,在batch size等于50,token length等于14的FP16數(shù)據(jù)格式下,整個Engram的耗時需要10.046ms, 而使用AMX來加速Linear轉(zhuǎn)換計算和Conv卷積計算,Engram的耗時只要6.022ms,整體性能提升至AVX-512的1.67倍ii。

圖4 測試得出的在FP16數(shù)據(jù)格式下分別用AMX 和AVX-512加速的執(zhí)行時間

(每次測試結(jié)果都可能存在少許浮動,在可接受范圍)

圖5 用英特爾® AMX加速Engram模塊中Linear和卷積計算的性能表現(xiàn)

上述探索和測試是基于DeepSeek開源的deepseek-ai/Engram: Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models的相關(guān)代碼,歡迎同行們一起討論交流。

接下來英特爾會繼續(xù)圍繞Engram及集成它的大模型開展更多、更為深入的探索,特別是在DeepSeek相關(guān)模型正式發(fā)布后,我們會更為全面、系統(tǒng)地評估CPU獨立運行和加速Engram模塊會對整個模型的端到端性能、部署和應用的門檻,以及投資回報等維度帶來怎樣的影響,相關(guān)進展與成果也將在第一時間分享。

正如開篇所說,我們相信這些工作與Engram的設(shè)計初衷是相向而行的,即在用 “查算分離” 理念打破傳統(tǒng)大模型的記憶困境的同時,也讓AI基礎(chǔ)設(shè)施中的GPU和CPU實現(xiàn)更好的協(xié)作,并充分釋放它們各自的潛能,進而大幅提升AI系統(tǒng)的部署效率及投資回報,或者進一步拉低AI部署與實踐的成本或門檻。我們的工作,就是希望能為這一目標的達成拓展出更為多樣化的技術(shù)路徑與更強的靈活性。