相鄰的 KV 緩存將合并為一個。在生成第一個字符時，KV 緩存長度為 1；生成第二個字符后，新生成的 KV 與前一個被合并，KV 緩存長度仍然保持為 1。這種動態(tài)合并機(jī)制有效壓縮了時間維度上的冗余信息。

然而，這也帶來了并行訓(xùn)練上的挑戰(zhàn)：雖然兩個時間步的 KV 緩存長度相同，但它們所包含的信息不同，若不加以區(qū)分，容易導(dǎo)致訓(xùn)練與推理行為不一致。

MTLA 通過一種優(yōu)雅的方式解決了這一問題。正如下圖所示，在訓(xùn)練階段，MTLA 保留了所有中間狀態(tài)的 KV 表達(dá)，并引入了步幅感知因果掩碼（stride-aware causal mask），確保每個 query 在訓(xùn)練時訪問到與推理階段一致的 KV 區(qū)域，從而準(zhǔn)確模擬增量推理中的注意力行為。

得益于這一設(shè)計，MTLA 能夠像標(biāo)準(zhǔn)注意力機(jī)制一樣通過矩陣乘法實(shí)現(xiàn)高效并行計算，在保持訓(xùn)練效率的同時完成對時間維度的壓縮。

此外，MTLA 還引入了解耦的旋轉(zhuǎn)位置編碼（decoupled RoPE）來建模位置信息，并對其進(jìn)行了時間維度上的壓縮，進(jìn)一步提升了整體效率。

值得強(qiáng)調(diào)的是，MTLA 不僅是一種更高效的自注意力機(jī)制，它還具備極強(qiáng)的靈活性與可調(diào)性。例如，當(dāng)將時間壓縮率 s 設(shè)置得足夠大時，MTLA 在推理過程中幾乎只保留一個 KV 緩存，這種形式本質(zhì)上就退化為一種線性序列建模方法。換句話說，線性序列建模可以被視為 MTLA 的極端情況，MTLA 在注意力機(jī)制與線性模型之間架起了一座橋梁。

然而，在許多復(fù)雜任務(wù)中，傳統(tǒng)注意力機(jī)制所具備的二次計算復(fù)雜度雖然代價高昂，卻提供了更強(qiáng)的建模能力。因此，MTLA 所引入的 “可調(diào)時間壓縮率 s” 這一設(shè)計思路，恰恰為模型提供了一個在效率與性能之間靈活權(quán)衡的可能空間。

MTLA 的卓越性能

MTLA 在一系列任務(wù)中展現(xiàn)了出色的性能，包括語音翻譯，文本摘要生成，語音識別和口語理解。例如在語音翻譯中，MTLA 在保持與標(biāo)準(zhǔn) MHA 相當(dāng)?shù)姆g質(zhì)量的同時，實(shí)現(xiàn)了超過 5 倍的推理速度提升，并將推理過程中的 GPU 顯存占用降低了超過 8 倍。

值得注意的是，僅當(dāng)時間壓縮率 s=2 時，MTLA 對 KV 緩存的壓縮程度就已經(jīng)與 MQA 相當(dāng)，且在模型性能上更具優(yōu)勢。而相比之下，MQA 所采用的減少 KV 頭數(shù)量的方法已達(dá)上限，而 MTLA 還有進(jìn)一步的空間。

未來發(fā)展

MTLA 具備在大規(guī)模場景中部署的顯著潛力，尤其是在大語言模型參數(shù)規(guī)模不斷擴(kuò)大、以及思維鏈等技術(shù)推動下生成序列日益增長的背景下，對 KV 緩存進(jìn)行時空壓縮正是緩解推理開銷的關(guān)鍵手段。在這樣的趨勢下，MTLA 有望成為未來大語言模型中自注意力模塊的重要替代方案。

當(dāng)然，與 DeepSeek 提出的 MLA 類似，MTLA 相較于 GQA 和 MQA，在工程落地方面的改動不再是簡單的一兩行代碼可以實(shí)現(xiàn)的優(yōu)化。這也意味著要將其大規(guī)模應(yīng)用到現(xiàn)有 LLM 框架中，還需要來自社區(qū)的持續(xù)推動與協(xié)同開發(fā)。

為促進(jìn)這一過程，MTLA 的實(shí)現(xiàn)代碼已全面開源，希望能夠?yàn)檠芯空吲c工程實(shí)踐者提供便利，共同推動高效注意力機(jī)制在大模型時代的落地與普及。