查算分離！用至強和AMX加速Engram，性能提升達1.67倍

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

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

Engram設(shè)計初衷:

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

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

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

Engram核心創(chuàng)新：

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

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

這種分工模式改變了傳統(tǒng)大模型“推理既要計算又要記憶”的狀態(tài)，就像給學(xué)者配備了一本可即時查閱的百科詞典，無需每次都從頭推導(dǎo)基礎(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”），每個子序列對應(yīng)其階數(shù)的多頭哈希機制，從預(yù)構(gòu)建的靜態(tài)N-gram 嵌入表（如圖中的2-Gram Embedding、3-Gram Embedding表）中以O(shè) (1)復(fù)雜度查詢對應(yīng)嵌入。

動態(tài)門控融合階段：N-gram 嵌入表先經(jīng)Concat（通道拼接）形成融合記憶向量，再通過兩個Linear（線性轉(zhuǎn)換）層分別投影為Key（記憶語義摘要）與Value（待注入信息）。隨后將當(dāng)前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)的復(fù)現(xiàn)、驗證及測試工作。我們的探索則更進一步——不同于原論文中Concat之后的工作任務(wù)將交還給GPU執(zhí)行，我們不僅將N-gram哈希檢索階段的計算放在至強® 平臺上執(zhí)行，還將動態(tài)門控融合階段中的Linear轉(zhuǎn)換計算和Conv卷積計算也放到該平臺上運行。換言之，我們是基于充分的性能調(diào)優(yōu)，嘗試用CPU平臺獨立運行和加速整個Engram模塊。

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

圖3 從Engram原有工作流程轉(zhuǎn)向用CPU運行整個模塊并導(dǎo)入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模塊會對整個模型的端到端性能、部署和應(yīng)用的門檻，以及投資回報等維度帶來怎樣的影響，相關(guān)進展與成果也將在第一時間分享。

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