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

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

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

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

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

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

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

MTLA 的卓越性能

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

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

未來發(fā)展

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

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

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