重點摘要：Mamba-3 是一款全新的狀態空間模型（SSM），其首要設計目標是「推論效率」——這與著重於訓練速度的 Mamba-2 截然不同。其關鍵升級在於更富表現力的遞迴公式、複數值狀態追蹤，以及能在不減緩解碼速度的前提下提升準確度的 MIMO（多輸入多輸出）變體。
結果顯示：在 1.5B 參數量級下，Mamba-3 SISO 在所有序列長度的預填（prefill）與解碼（decode）延遲表現上，均超越了 Mamba-2、Gated DeltaNet，甚至勝過了 Transformer 架構的 Llama-3.2-1B。
研究團隊已開源其核心算子（kernels），這些算子混合使用了 Triton、TileLang 和 CuTe DSL 編寫，以最大化硬體效能。
本文同步發布於 Goomba Lab 網誌，內容涵蓋卡內基美隆大學、普林斯頓大學、Cartesia AI 與 Together AI 研究人員的合作成果。
自 2024 年中 Mamba-2 發布以來，大多數架構已從 Mamba-1 轉向。原因為何？Mamba-2 押注於「訓練效率」是狀態空間模型（SSM）的最大瓶頸，因此簡化了底層 SSM 機制，使其訓練速度比前代快上 2 至 8 倍，進而獲得更廣泛的採用。
此後，大型語言模型（LLM）的版圖開始轉變。雖然預訓練仍然至關重要，但更多關注點已轉向微調後（post-training）與部署階段，這兩者都極度依賴推論效能。特別是隨著可驗證獎勵的強化學習（RLVR）應用於程式碼或數學領域，需要海量的生成式滾動（rollouts）；而近期如 Codex、Claude Code 甚至 OpenClaw 等代理（Agentic）工作流的興起，更是將推論需求推向了頂峰。
儘管推論的重要性日益顯著，許多線性架構（包含 Mamba-2）仍是從「訓練優先」的角度出發開發的。為了加速預訓練，底層 SSM 被逐步簡化（例如將對角過渡矩陣簡化為純量乘以單位矩陣）。雖然這帶來了訓練速度的提升，卻使得推論步驟變得「過於簡單」且完全受限於記憶體頻寬——GPU 的核心運算能力並未充分發揮，大部分時間都花在搬移記憶體資料上。
在這個推論為王的新时代，我們非常關注如何突破品質與效率的邊界：我們希望更好的模型能跑得更快。
一個自然的問題浮現：
如果以推論為核心設計 SSM，它會長什麼樣子？
Mamba-3 模型
缺了什麼？線性模型的主要吸引力在於其名：由於狀態大小固定，運算量隨序列長度呈線性增長。然而，天下沒有白吃的午餐。同樣的固定狀態大小雖然實現了高效運算，卻強迫模型將所有過去資訊壓縮成單一表示，這與 Transformer 透過不斷增長的狀態（KV cache）來儲存所有過去資訊的做法形成根本上的對比。因此，如果我們無法擴大狀態，該如何讓這個固定狀態做更多工作？
我們觀察到，先前的設計簡化了遞迴與過渡矩陣以加速訓練。然而，這項改變也降低了動態的豐富度，並使解碼受限於記憶體：每個 token 的更新所進行的運算量，相對於記憶體搬移量來說實在太少。這為我們提供了三個著力點：(1) 讓遞迴本身更具表現力，(2) 使用更豐富的過渡矩陣，以及 (3) 在每次更新中加入更多平行且幾乎免費的運算工作。
基於這些見解，我們在三個核心方面改進了 Mamba-2：
- 透過我們提出的指數梯形離散化方案（exponential-trapezoidal discretization scheme），推導出更通用的遞迴公式，提升 SSM 機制的表現力。
- 透過建模複數值 SSM 系統，擴展狀態追蹤能力。
- 採用多輸入多輸出（MIMO）SSM，平行建模多個 SSM，取代現行的單輸入單輸出（SISO）SSM，從而在對解碼延遲影響極小的情況下提升整體效能。
透過這三項變革，Mamba-3 在維持相似推論延遲的同時，推動了效能邊界。
值得注意的是，這三項變革皆啟發自更「經典」的控制理論與狀態空間模型文獻。
我們的工作逆勢而為，不同於許多現代線性架構採用替代性的遞迴詮釋（如線性注意力或測試時訓練），那些方法不易捕捉這些概念。
架構
Mamba-2 層級有何變化？除了上述三個核心 SSM 的方法論升級外，我們還調整了架構，使其更符合現代語言模型的慣例。
根據圖示，你會注意到我們做了一些更動。概略來說：
正規化（Norms）：我們加入了 QKNorm，或在 SSM 術語中稱為「BCNorm」，實證顯示這能穩定 Mamba-3 模型的訓練。加入此正規化使 Mamba-3 與當前的 Transformer 及 Gated DeltaNet (GDN) 模型看齊。有了 QKNorm，Mamba-2 中的 RMSNorm 便成為可選項。然而，實證發現，在混合模型中保留它可能仍有價值，因為它有助於長度外推能力。稍後會再詳述。
告別短卷積（Goodbye Short Conv）：我們成功移除了 Mamba-1/2 中令人頭痛的短因果卷積（short causal convolution），方法是結合 (1) BCNorm 後對 B 和 C 矩陣加上簡單偏置，以及 (2) 我們新的基於離散化的遞迴。新的遞迴隱式地對輸入至隱藏狀態應用了卷積，我們將在文章第二部分展示這是如何實現的。
真的可以移除短卷積嗎？
Mamba-3 的改變在SSM 遞迴內部加入了類卷積組件，但這與放置在SSM 遞迴外部的標準短卷積並不完全互換。
後者仍可與 Mamba-3 搭配使用，但我們基於實證決定不採用。研究發現，加回標準短卷積：
- 並不會提升效能；事實上，它會略微降低效能。
- 在更多真實世界任務（如 NIAH）上，並不會削弱檢索能力。話雖如此，若無短卷積，在小型合成任務（如 MQAR）上的訓練會變得較困難。但由於真實世界的檢索行為未受影響，我們不認為這是主要限制。
至於原因？我們未深入研究理論機制，但在論文中我們假設，BC 偏置與指數梯形遞迴都執行了類似的類卷積機制，在實證上發揮了與外部短卷積相同的功能。
短卷積簡史
短卷積現已是多數高效能線性模型的核心組件。短卷積的變體最早由 H3（以「移位 SSM」形式，靈感來自 Anthropic 的「smeared」歸納頭研究）和 RWKV-4（透過其「token shift」機制）應用於遞迴架構，而後由 Mamba-1 以其當前形式普及。
它之所以如此普遍，是因為先前的研究反覆證明短卷積能提升實證效能，並在理論上支持歸納式檢索能力。
最後，你會注意到幾個新組件， namely RoPE 和 MIMO 投影。RoPE 模組透過將複數過渡詮釋為旋轉，來表達複數值 SSM，避免了昂貴的核心重實作。MIMO 投影則將 B 和 C 矩陣擴展至 MIMO SSM 所需的適當表示形式。
我們將在文章第二部分深入探討這兩點的動機與確切實作方式（那裡有很多乾貨 🎁），所以目前只需將它們視為獨立且基礎性的改進，各自為模型的效能與能力做出貢獻。
最後，我們的整體架構現在遵循 Transformer 與其他線性模型的標準慣例，採用交錯的 MLP 層。
實證結果
我們將最終的 Mamba-3 模型與其他流行的線性替代方案及 Transformer 基準進行評估。
語言建模
我們發現，新的 Mamba-3 模型在各個預訓練模型規模的語言建模任務上，優於先前的 Mamba-2 模型以及強大的線性注意力替代方案（如 GDN）。Mamba-3-SISO 可直接與先前的線性模型相比；例如，它在架構形狀（模型維度、狀態大小等）上與 Mamba-2 完全匹配，且訓練時間相當。我們的 Mamba-3 MIMO 變體在 1B 規模下，將下游任務的準確度進一步提升了 1 個百分點以上，需要注意的是，MIMO 需要更長的訓練時間，但不會增加解碼延遲！
為何訓練成本上升但推論沒有？
雖然我們將在文章第二部分詳細討論，但這裡先給讀者一個預告：
這種二分法可追溯至訓練與推論分別受限於運算或記憶的本質。當前的線性模型被設計為大量使用GPU 張量核心（Mamba-2 的主要貢獻之一）以進行快速訓練，但在解碼期間，每個時間步所需的運算量極少，導致硬體大部分時間處於閒置狀態。
因此，如果我們圍繞著增加每個時間步所需的 FLOPs 來設計架構，推論延遲將大致保持不變，因為我們只需利用部分閒置核心即可——但這對訓練來說可不行！
檢索任務
線性模型因其固定大小的狀態，在基於檢索的任務上自然表現較差。正如預期，在純模型中，Transformer 在檢索任務上表現優異，但 Mamba-3 在次二次方（sub-quadratic）替代方案中表現良好。有趣的是，加入 MIMO 進一步提升了檢索效能，且未增加狀態大小。
鑑於這種先天劣勢但整體建模效能強勁，
我們預測，未來線性層將主要與全域自注意力層結合使用。*
*至少就語言建模而言
混合模型結合了線性層類似的記憶體般特性與自注意力 KV 快取精確的資料庫般儲存能力，實證顯示其表現優於純模型，同時能顯著節省記憶體與運算。我們也發現，線性層與自注意力的結合，其檢索能力優於標準 Transformer。
然而，我們必須強調，這些線性模型與自注意力相互作用的確切方式尚未被完全理解。例如，我們發現使用 Mamba-3 的可選預輸出投影（pre-output projection）能提升合成 NIAH 任務上的長度泛化效能，但代價是輕微降低了情境內真實世界檢索任務的表現。此外，即使是回傳正規化的細節，如位置（閘道前 vs 閘道後）與類型（分組 vs 常規），都會對由半結構化和非結構化數據組成的任務（如 FDA 和 SWDE）的準確度產生不可忽視的影響。
無所不在的核心（Kernels）
我們很期待看到大家用 Mamba-3 建構出什麼。為了促進這一點，我們將開源我們的算子（kernels），其速度與原始的 Mamba-2 Triton 算子相當。
延遲基準測試
預填（Prefill）延遲
預填 + 解碼（Prefill+decode）延遲
以上為 1.5B 模型在單一 H100-SXM 80GB GPU 上，不同序列長度的預填與預填 + 解碼（兩者 token 數量相同）延遲。所有序列長度均使用 128 的批次大小，報告的是三次重複測試的掛鐘時間（秒）。
在 1.5B 規模下比較各模型時，Mamba-3（SISO 變體）在所有序列長度上實現了最快的預填 + 解碼延遲，超越了 Mamba-2、Gated DeltaNet，甚至超越了擁有高度優化 vLLM 生態系統的 Transformer。此外，Mamba-3 MIMO 在速度上與 Mamba-2 相當，但效能更強。
Mamba-3 SISO 基於 Triton 的預填保持了與 Mamba-2 幾乎相同的效能，證明新的離散化和數據依賴的 RoPE 嵌入並未引入額外開銷；而 Mamba-3 MIMO 僅因其高效的 TileLang 實作而在預填上產生了适度的減速。兩種 Mamba-3 變體強勁的解碼效能，部分歸功於 CuTe DSL 的實作，這得益於 Mamba-3 組件的簡潔性而變得更加容易。
設計選擇
我們花了大量時間思考如何在妥協易用性的前提下，將算子速度極致化。我們最終採用了以下技術棧：Triton、TileLang 和 CuTe DSL。
選擇 Triton 相當容易。它幾乎是架構開發的標準（優秀的 flash linear attention 倉庫完全使用 PyTorch 和 Triton 構建），這是有原因的，因為它透過受控的分塊（tiling）和算子融合，實現了比標準 PyTorch 更好的效能，同時又是一種平台無關的語言。Triton 還有一些相當巧妙的功能，例如 PTX（一種面向 GPU 的組合語言）注入及其對張量記憶體加速器（TMA，在 Hopper GPU 上）的支持，可實現從全域記憶體到共享記憶體的大量非同步傳輸。
我們的 MIMO 預填算子則是使用 TileLang 開發的。對應該變體的額外投影提供了一個機會，讓我們可以透過策略性地操作 GPU 記憶體層級結構來減少記憶體 IO。不幸的是，Triton 無法提供我們所需的記憶體控制粒度，因此我們選擇了 TileLang，它允許我們明確聲明和控制共享記憶體區塊（tiles）並創建暫存器片段，從而更有效地重用記憶體，同時仍保持足夠高階，讓我們能快速開發算子。
既然我們一直強調推論與解碼的重要性，我們決定為解碼算子使用 CuTe DSL。透過其 Python 介面，我們能夠使用來自 CUTLASS 的高階抽象來生成低階算子。在這裡，我們實際上擁有 CUDA 等級的控制權，使我們能夠開發出針對硬體規格（此處為 Hopper GPU）量身打造的高效能算子。透過對張量佈局和 warp 專門化的細粒度控制，我們構建了一個能充分利用 GPU 所有功能的算子。
重要的是，這些跨越不同 GPU 抽象層次的實作之所以成為可能，歸功於 Mamba-3 簡單、輕量級新增內容及其巧妙實例化的底層演算法設計。我們在完整發布版中更深入地討論了確切的融合結構和算子 DSL 等細節。
接下來
很高興你讀到了第一部分的結尾！關於我們的算子、實驗結果和消融研究，還有許多細節來不及在這篇文章中介紹，但別擔心！所有內容都能在 我們的論文 中找到，算子已開源於 mamba-ssm！
接下來，系列作的 第二（也是最後）部分 將深入探討 Mamba-3 的三項核心改進及其 SSM 基礎，並提出我們特別感興趣的方向。
參考文獻
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces [PDF]
Gu, A. and Dao, T., 2024. - Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality [PDF]
Dao, T. and Gu, A., 2024. - Gated Delta Networks: Improving Mamba2 with Delta Rule [PDF]
Yang, S., Kautz, J. and Hatamizadeh, A., 2025. - Learning to (Learn at Test Time): RNNs with Expressive Hidden States [PDF]
Sun, Y., Li, X., Dalal, K., Xu, J., Vikram, A., Zhang, G., Dubois, Y., Chen, X., Wang, X., Koyejo, S., Hashimoto, T. and Guestrin, C., 2025. - Hungry Hungry Hippos: Towards Language Modeling with State Space Models [PDF]
Fu, D.Y., Dao, T., Saab, K.K., Thomas, A.W., Rudra, A. and Ré, C., 2023. - In-context Learning and Induction Heads
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S. and Olah, C., 2022. Transformer Circuits Thread. - RWKV: Reinventing RNNs for the Transformer Era [PDF]
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Grella, M., GV, K.K., He, X., Hou, H., Lin, J., Kazienko, P., Kocon, J., Kong, J., Koptyra, B., Lau, H., Mantri, K.S.I., Mom, F., Saito, A., Song, G., Tang, X., Wang, B., Wind, J.S., Wozniak, S., Zhang, R., Zhang, Z., Zhao, Q., Zhou, P., Zhou, Q., Zhu, J. and Zhu, R., 2023. - Test-time regression: a unifying framework for designing sequence models with associative memory [PDF]
Wang, K.A., Shi, J. and Fox, E.B., 2025. - An Empirical Study of Mamba-based Language Models [PDF]
Waleffe, R., Byeon, W., Riach, D., Norick, B., Korthikanti, V., Dao, T., Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., Kulshreshtha, G., Singh, V., Casper, J., Kautz, J., Shoeybi, M. and Catanzaro, B., 2024.