查算分離！用至強(qiáng)和AMX加速Engram，性能提升達(dá)1.67倍

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

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

Engram設(shè)計(jì)初衷:

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

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

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

Engram核心創(chuàng)新：

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

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

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

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

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

Engram計(jì)算流程解析

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

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

N-gram哈希檢索階段：模塊對(duì)輸入序列提取2-gram、3-gram等不同子序列（如圖中的“the Great”、“Alexander the Great”），每個(gè)子序列對(duì)應(yīng)其階數(shù)的多頭哈希機(jī)制，從預(yù)構(gòu)建的靜態(tài)N-gram 嵌入表（如圖中的2-Gram Embedding、3-Gram Embedding表）中以O(shè) (1)復(fù)雜度查詢對(duì)應(yīng)嵌入。

動(dòng)態(tài)門控融合階段：N-gram 嵌入表先經(jīng)Concat（通道拼接）形成融合記憶向量，再通過(guò)兩個(gè)Linear（線性轉(zhuǎn)換）層分別投影為Key（記憶語(yǔ)義摘要）與Value（待注入信息）。隨后將當(dāng)前Transformer層的Input Hidden（全局上下文）與 Key 做 Scaled Dot Product （點(diǎn)積計(jì)算）生成門控權(quán)重，動(dòng)態(tài)過(guò)濾與上下文無(wú)關(guān)的Value信息。加權(quán)后的Value再通過(guò)Conv卷積計(jì)算完成局部融合。

殘差集成階段：局部融合后的記憶特征通過(guò)加法操作，殘差加回該Transformer Block的輸入，并直接輸入后續(xù)的Attention與MoE層。

英特爾的探索：

用CPU獨(dú)立運(yùn)行及加速整個(gè)Engram模塊

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

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

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

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

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

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

（每次測(cè)試結(jié)果都可能存在少許浮動(dòng)，在可接受范圍）

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

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

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

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