博客 / 詳情

返回

阿里雲 Tair 聯合 SGLang對 Mamba-Transformer 等混合架構模型的支持方案

導讀

接着上一節內容對KV Cache存儲方案的深入解讀,本文介紹了阿里雲 Tair KVCache 團隊與SGLang 社區在推理框架上的提效——支持混合架構模型的工程化實踐。
在大模型長文本與智能體化趨勢下,Transformer 面臨顯存與計算瓶頸,而高效的 Mamba 模型語義召回受限。混合架構通過結合兩者優勢應運而生,卻帶來系統級挑戰:Transformer 的 Token 粒度緩存與 Mamba 的請求粒度原地更新機制存在本質衝突,導致前綴緩存、推測解碼等傳統優化技術失效。這迫切要求推理框架進行架構重構,以解決異構狀態管理與調度的難題。
本文在 SGLang Hybrid Models 的工作基礎上,深入剖析混合架構的設計原理、實現難點與系統級優化路徑,為高效、可靠的大模型混合推理提供可落地的技術方案。
混合架構:SGLang 首創了雙內存池,完美兼容 Transformer 和 Mamba 兩種截然不同的內存習性。
技術方案:通過狀態快照技術,解決了 Mamba 模型“無法回滾”的缺陷,讓緩存複用和推測解碼成為可能。
優化效果:實測 Qwen3-Next 等混合模型在 SGLang 上跑得飛快。
本系列技術文章將系統性拆解面向智能體推理的 KVCache 技術演進路徑:

  1. 智能體式推理對 KVCache 的挑戰與 SGLang HiCache 技術深度剖析
  2. 3FS-KVCache 工程化落地:企業級部署、高可用運維與性能調優實踐
  3. 本文 | Hybrid Model Support:SGLang 對 Mamba-Transformer 等混合架構模型的支持方案
  4. Tair KVCache Manager:企業級全局 KVCache 管理服務的架構設計與實現
  5. KVCache 仿真分析:高精度的計算和緩存模擬設計與實現
  6. Hierarchical Sparse Attention:分層稀疏注意力框架下的 KV 分層管理與按需加載
  7. 展望:KVCache驅動的軟硬結合演進

Tair KVCache 作為阿里雲數據庫Tair產品能力的延伸,本質是緩存範式的三次躍遷:
🔹 從 Redis 的 “緩存數據 → 減少 I/O”;
🔹 到 GPU KVCache 的 “緩存計算中間態 → 減少重複計算”;
🔹 再到 Tair KVCache 的 “規模化、智能化的注意力狀態管理 → 重構大模型推理成本模型”它標誌着緩存正從輔助組件升級為 AI 基礎設施層的核心能力——讓“狀態”可存儲、可共享、可調度,支撐智能體時代的規模化推理底座。

1. 引言

1.1 混合架構的崛起

在大語言模型推理服務邁向長上下文、多模態交互與智能體化的新階段,傳統架構的侷限性日益凸顯。Transformer 模型憑藉其注意力機制在語義建模上表現卓越,但其計算開銷隨序列長度呈平方級增長,KVCache內存佔用線性膨脹,其在超長文本、持續對話等場景下面臨顯存限制與算力瓶頸。與此同時,以Mamba 為代表的狀態空間模型通過線性計算複雜度和恆定的內存消耗開闢了新路徑,但其有限的狀態容量與不可逆的上下文壓縮機制,又難以支撐複雜推理任務所需的細粒度語義召回能力。
這一矛盾催生了混合架構的崛起——將 Transformer 的全注意力層與 Mamba 的狀態空間模型層交錯設計,試圖在效率與性能間尋求平衡點。然而,混合模型的落地並非簡單的模塊堆砌,其背後隱藏着更深層的系統級挑戰。本文在SGLang Hybrid Models[1]的工作基礎上深入剖析其設計原理、實現難點與優化路徑,為基於混合架構的高效LLM 推理架構提供實踐參考。

1.2 狀態空間模型:線性效率與有限容量的權衡

狀態空間模型(State Space Models, SSMs),通過遞歸式上下文壓縮技術,將動態變化的token序列映射為固定維度的隱式狀態。這種設計在計算範式上實現了雙重突破:
1)內存效率提升:推理過程中狀態維度恆定(image),擺脱傳統注意力機制隨序列長度線性膨脹(image)的內存瓶頸;
2)計算複雜度降低:自迴歸生成時計算量僅隨序列長度線性增長(image),相較注意力機制的平方級複雜度(image)實現數量級優化。
1

然而,這種設計存在潛在約束:有限的狀態容量與不可逆的壓縮機制。SSM的固定維度狀態如同“信息漏斗”,在長程上下文建模中難以保留細粒度的敏感特徵,導致複雜模式匹配與精確語義召回能力顯著弱於注意力架構。這一缺陷在需要多跳推理、長文檔分析等場景尤為突出,成為制約純Mamba模型落地的難題。
為突破這一困境,混合架構應運而生——通過設計注意力與SSM層間交錯的模型,將SSM的線性效率與注意力的語義建模能力深度融合。以Qwen3-Next、Kimi-Linear為代表的先進模型採用注意力層與SSM層混合配比的架構,在長上下文任務中實現雙重增益:通過全注意力層維持對關鍵語義特徵的捕捉能力,高效地保留長上下文推理能力;SSM層替代部分注意力計算,顯著降低內存帶寬壓力與計算延遲,提升吞吐效率。

1.3 當前系統的挑戰

由於注意力層與SSM層在計算範式存在根本性差異,混合架構模型的工程化落地需完善考慮不同類型層間的狀態管理和系統級優化實現。
首先,需要解決注意力層與SSM層不同計算範式的資源協同調度難題:注意力層依賴前序KVCache進行計算,SSM層則依賴固定維度的SSM狀態進行推理。兩者計算範式的區別帶來內存管理的差異:注意力層運行時依賴token粒度的KVCache管理,而SSM層則可以以請求粒度維護SSM狀態。這種差異給推理系統管理混合架構模型的KVCache與SSM狀態帶來挑戰。
注意力層與SSM層狀態管理機制的不一致提升了推理優化策略的適配難度。SSM層會“原地覆蓋式”的更新狀態,這種壓縮特性形成不可逆的更新路徑,這與智能體場景下的前綴緩存、分支推測解碼等需要狀態回滾的優化策略產生衝突。當系統嘗試複用跨請求的共享上下文時(如多用户共用的系統指令模板或知識庫文檔),傳統基於KVCache塊的空間共享機制因無法兼容SSM狀態的原地更新特徵而失效。系統需要設計跨注意力KVCache和狀態空間模型SSM不同模式的聯合緩存協議,這種跨層狀態同步不僅需要考慮內存管理複雜度,還需要解決潛在的競態條件。
以前綴緩存為例,假設我們需要基於文檔1回答兩個問題。在注意力場景中,由於KVCache是以token粒度維護,在回答問題1時文檔1的KVCache便自然地以token粒度計算維護好,當我們希望回答問題2時可以直接複用文檔1的KVCache。而在狀態空間模型場景,SSM狀態會被原地式覆蓋,如果不顯式地在推理過程中將某個時間點的SSM狀態緩存下來,當問題1回答完成時,系統只會保留完成問題1回答後的SSM狀態SSM p+3,文檔1的完成計算狀態SSM n是缺失的。此時問題2的回答就需要重頭開始計算,前綴緩存失效。
2
在分佈式部署層面,當前主流的PD分離架構以及KVCache的多層存儲體系,均圍繞注意力機制的計算特性進行了深度優化。KVCache通常以token或page為粒度,在SGLang推理實例之間,或在SGLang與底層存儲引擎之間實現高效傳輸與共享,從而在用户體驗上保障更嚴格的SLA,在推理性能上支持上下文複用等“以存代算”的優化策略。如何在現有分佈式推理框架中擴展緩存與通信機制,使其既能保留對注意力層KVCache的高效支持,又能兼容SSM層中SSM的狀態緩存、跨節點傳輸與持久化能力成為推動此類模型工程化落地的關鍵挑戰。

2. 內存管理

2.1 雙池內存架構

為應對混合架構模型在內存管理方面帶來的獨特挑戰,SGLang提出了多池內存架構。該設計的核心理念在於:深入識別不同注意力機制組件所表現出的差異化內存行為特徵,並據此制定針對性強、精細化的內存管理策略。
具體而言,在SGLang框架中,傳統注意力層生成的KV Cache表現出“細粒度增長、短週期波動”的特性——每個新生成的token僅產生數KB級別的緩存數據,並隨着推理過程動態累積與釋放。相比之下,混合架構中新引入的狀態空間模型機制依賴的SSM狀態則呈現出“大塊連續、長週期持有”的特點:單個請求所需的SSM狀態通常佔用數MB的存儲空間,且必須完整保留直至該請求完全結束。若將這兩種內存需求差異顯著的數據結構混置於同一內存池中,不僅會因大小懸殊(KB 級 vs. MB 級)的分配單元交替出現引發嚴重的內存碎片問題,還會顯著增加系統實現的工程複雜度與運行時開銷。
為此,SGLang採用物理隔離的雙內存池設計,將整體內存劃分為兩個固定大小的獨立區域:狀態空間模型Mamba 狀態池 和 注意力KV Cache池。兩者的總容量在服務啓動時即通過 --mamba-full-memory-ratio參數靜態配置並預分配,從而有效規避了運行時動態分配可能引發的OOM風險。
其中,Mamba狀態池以請求為單位進行管理:藉助HybridReqToTokenPool數據結構,系統在請求初始化階段即為其分配一個固定大小(通常為MB級)的連續內存頁,並將其生命週期與請求綁定,請求完成後立即回收,確保高效利用大塊內存。而KV Cache池則延續細粒度管理策略,通過HybridLinearKVPool實現注意力層與物理內存的映射,專用於支持全注意力計算。這種分離式設計不僅避免了在SSM層中分配無效KV Cache,還實現了兩類內存需求的正交管理,顯著提升了整體內存利用率。
3

2.2 彈性內存池

然而,固定比例的池劃分難以適應真實場景中波動的工作負載。例如,當系統負載從高併發的短對話任務切換至低併發但上下文極長的生成任務時,Mamba池往往因請求減少而閒置,而KV Cache池卻因長序列緩存需求激增而迅速耗盡,進而限制批處理規模引發性能瓶頸。
為此,SGLang在多池隔離架構的基礎上引入了彈性內存池機制,在保持Mamba狀態與KV緩存語義隔離的前提下,實現池間容量的運行時動態重分配。該機制首先在存儲管理層依託CUDA虛擬內存管理能力:系統在啓動時預分配一個超額預定的虛擬地址空間,併為每個內存池創建可彈性伸縮的張量數據結構。這些張量本身不立即佔用物理顯存,而是作為虛擬佔位符。當某類緩存需求增長時,控制模塊將物理顯存頁動態映射到對應的虛擬地址區間,實現“按需激活”的內存分配;反之,當某內存池使用率下降,其空閒塊所佔的物理頁會被解除映射並釋放,從而回收資源。以長文本生成為例,當負載由短請求轉為長序列任務時,推理批大小通常減小,SSM層所需的SSM狀態總量隨之降低,Mamba池使用率下降,系統便可自動將其空閒物理頁轉移至KV Cache池,支持更長上下文的持續擴展,有效緩解靜態分配導致的內存利用率不均問題。
在控制決策層面,系統通過一個集中式調度模塊實現智能、安全的池間資源再分配。各內存池在初始化階段向該模塊註冊元信息。運行時,若某一池因容量不足發起擴容請求,控制模塊會實時評估所有池的當前使用率,選擇最空閒的池觸發縮容操作——即釋放其部分物理顯存頁,並在確認釋放成功後,授權請求方完成擴容。整個過程嚴格限定在固定的總GPU顯存預算內,無需重啓服務或重新分配全局內存,既避免了 OOM風險,又保障了分配操作的原子性與安全性。
4

通過“多池隔離 + 彈性調度”的協同設計,SGLang既保留了針對不同內存訪問模式(大塊連續 vs. 細粒度動態)的精細化管理優勢,又具備對動態工作負載的自適應能力,在保障系統穩定性的同時提升了GPU顯存的整體利用效率,為更大批次或更長上下文的高效推理提供了堅實支撐。

3. 關鍵技術優化與適配

3.1 混合前綴緩存

在語言模型推理優化領域,前綴緩存通過複用不同請求之間的公共前綴計算結果,顯著提升系統吞吐與效率。然而,當該技術應用於融合了狀態空間的混合架構時,會遭遇一系列挑戰。全注意力層的前綴緩存依賴於KVCache的token粒度管理,可基於前綴匹配截斷,而SSM層中的SSM狀態管理機制則呈現出截然不同的特性:其狀態在推理過程中採用原地更新方式,無法像全注意力層的KVCache 那樣通過簡單截斷序列實現狀態回滾,因而難以精確還原任意歷史前綴對應的狀態;同時,單個SSM狀態通常達MB量級,相較於單個token的KVCache以數量級的形式增長,token粒度的緩存會導致存儲開銷急劇上升;更關鍵的是,大多數SSM狀態緩存具有“全有或全無”的複用特性——一個SSM狀態緩存只有當計算它的前綴全部匹配時才能被複用,不支持部分或增量式狀態複用。這些因素導致難以將傳統 Radix 樹結構用於此類混合模型。
為應對上述挑戰,SGLang引入了新的的Radix樹MambaRadixCache——一種專為混合狀態空間模型和注意力模型設計的混合前綴樹結構。該數據結構在不用修改已有Mamba推理算子的前提下,實現了對Mamba狀態與KVCache緩存的協同高效管理。
在匹配階段,系統在Radix樹中查找與當前輸入具有最長公共前綴且已緩存有效SSM狀態的節點。KVCache緩存由於其一經寫入不會修改的不變性,可以直接引用匹配上的KVCache進行復用。SSM狀態則會在後續推理時原地更新,需要將匹配的狀態完整拷貝快照給新請求,以避免多個併發請求因共享狀態導致的相互干擾,確保狀態隔離性與推理正確性。
在插入階段,系統在完成Chunked Prefill或逐token解碼後,將KVCache緩存與SSM狀態分別寫入Radix樹:KVCache 緩存僅需記錄對應內存頁的索引,而SSM狀態則需分配新的內存頁進行狀態拷貝,並將新頁索引關聯至相應樹節點。
在驅逐階段,MambaRadixCache採用雙LRU隊列機制,分別追蹤KV緩存與SSM狀態的訪問時間戳。其中KV緩存的驅逐嚴格遵循從葉節點向根節點逐層回收的原則,以維護Radix樹拓撲結構的完整性,而SSM狀態則採用更靈活的彈性驅逐策略,允許從任意層級節點釋放內存。
5

通過這一設計,MambaRadixCache幫助混合了SSM和注意力層的模型能夠在無需修改任何底層算子或Mamba推理算子的前提下無縫集成高效的前綴緩存能力。該方案不僅保持了原始計算邏輯的簡潔與高性能,還顯著降低了重複計算開銷與內存佔用,為大規模高效推理提供了關鍵基礎設施支持。

3.2 推測解碼適配方案

推測解碼作為大模型推理加速的核心技術,在全注意力架構中通過並行生成並驗證候選Token序列,顯著提升了推理效率。然而,當將其應用於狀態空間模型時,卻面臨根本性的適配挑戰。其根源在於SSM的狀態更新機制與傳統注意力中的KV Cache存在本質差異:SSM採用原地更新策略,每處理一個新 Token,其內部狀態image更新可以簡單抽象為遞推公式:image,會被不可逆地覆蓋。這種設計雖然在序列建模中高效簡潔,卻使得系統在推測解碼的驗證階段無法像處理KV Cache那樣簡單截斷或回滾——一旦某個候選Token被拒絕,其對SSM狀態的修改已永久生效,歷史狀態無法恢復。

更進一步,現有推測解碼方法如Eagle-Tree所依賴的注意力掩碼機制,也與SSM的狀態演化邏輯不兼容。Eagle-Tree 通過動態構建注意力掩碼來支持多路徑並行驗證,而SSM並不顯式維護Token間的注意力關係,其狀態是全局累積無局部掩碼控制的,無法直接適用。

為應對這些挑戰,SGLang提出了一種基於緩存隔離的新架構:為每個候選Token分配獨立的Mamba緩存槽,從而構建物理隔離的狀態沙箱。以三級候選序列 “the → air → streets” 為例,系統會分別在三個緩存槽中維護遞進的狀態演化——槽 1 存儲基礎狀態經 “the” 更新後的結果,槽 2 在此基礎上注入 “air”,槽 3 則繼承前一狀態並加入 “streets”。當驗證器確認 “the streets are” 這一前綴有效後,無需重新計算中間步驟,只需將對應槽(如槽 3)中的最終狀態直接提升為主SSM狀態,實現高效、無損的狀態切換。

在更復雜的 Top-K > 1 場景下每步會生成多個候選分支,該方案進一步引入父節點索引預計算機制。在推測生成階段,系統為每個候選 Token 顯式記錄其在推測樹中的父節點;進入驗證階段後,依據該索引追溯至對應的父狀態,並執行遞歸更新image。這一設計不僅保留了Eagle-Tree的多路徑探索能力,還使其與SSM的狀態演化機制對齊,成功將高效的推測解碼擴展至SSM架構,為其實時推理提供了可行路徑。
6

3.3 PD 分離架構擴展

SGLang的PD分離架構通過擴展傳輸協議,引入了面向不同注意力機制的專用狀態傳輸通道,高效支持混合模型的分離式部署。在標準的分頁KVCache傳輸之外,系統還為各類模型特有的內部狀態——例如Mamba中的SSM狀態、滑動窗口注意力中的窗口緩存等——設計了獨立的並行數據路徑,實現對非注意力KVCache狀態的高效傳輸。這種設計使得系統能夠靈活適配多種新型注意力機制,而無需對核心調度和通信邏輯進行大規模重構。
以同時包含注意力層和SSM層的混合模型為例,系統維護兩個相互獨立的內存池:一個用於注意力層的分頁 KVCache 池,另一個專用於存儲 SSM層所需的SSM狀態的Mamba狀態池。當新請求到達時,系統首先通過 MambaRadixTree 進行前綴匹配;若命中緩存,則將匹配到的MambaState複製至為該請求新分配的Mamba狀態緩衝區,並以此為基礎繼續執行Prefill推理。Prefill完成後,Prefill實例會將最終的Mamba狀態作為一個連續的內存塊,以原子的方式一次性傳輸至Decode實例,後者通過 dst_state_indices 告知Prefil實例接收該狀態的目標槽位。與支持增量傳輸的分頁KV Cache不同,Mamba狀態必須整體傳輸,無法分段發送。為確保狀態正確就位,Decode實例在請求調度階段即預先分配好對應的KV Cache頁面槽位和專用的Mamba狀態槽位,使接收到的狀態能夠準確寫入後續Decode步驟所需的內存位置,從而保障推理的連續性與正確性。
若要在現有PD架構中集成一種新的混合狀態池以支持分離式服務部署,僅需在當前實現基礎上完成少量擴展。首先,暴露該狀態類型所對應的緩衝區指針、總大小及單個條目長度,以便將其註冊到統一的傳輸系統中。其次,在Prefill和Decode工作節點中分別定義state_indices的生成邏輯,明確待傳輸狀態的源地址與目標地址;這一邏輯需根據注意力機制的特性進行定製——例如,全注意力或稀疏注意力層通常使用token或block粒度的KV Cache頁索引,SSM層採用請求粒度的單一索引,而滑動窗口注意力則可基於窗口頁索引進行管理。最後,為該狀態類型在KV Cache管理器中註冊一個唯一的state_type標識符,並在後端傳輸模塊中添加對應的讀寫、傳輸處理邏輯。整個過程高度模塊化,無需侵入核心調度流程。
7

通過上述機制,SGLang實現了對異構模型狀態的統一、高效且可擴展的管理,不僅兼容傳統Transformer架構,也能無縫支持Mamba、SWA等新興注意力變體,為混合架構大模型的高性能分離式推理提供了堅實基礎。

4. 性能驗證

SGLang在v0.5.5版本用H200跑Qwen3-Next-80B-A3B-Instruct-FP8的實驗驗證了上述設計的有效性。如下圖所示,啓用前綴匹配以存代算的能力可以避免重複計算匹配的前綴,將TTFT降低至57.63%
8

對 Qwen3-Next-80B-A3B-Instruct-FP8 模型在批量大小(batch size)為 1 的條件下進行了推測解碼的性能測試:
● 當 MTP 窗口大小為 2 個 token、top-k=1 時,系統吞吐量達到 257.20 tokens/秒,平均接受長度為 2.709 個 token。
● 當 MTP 窗口擴大至 3 個 token、top-k 仍為 1 時,吞吐量提升至 306.94 tokens/秒,平均接受長度增至 3.413 個 token。
● 進一步將 MTP 窗口設為 4 個 token,並採用 top-k=4 及 8 個draft token的配置,吞吐量進一步提升至 324.57 tokens/秒,平均接受長度達到 4.231 個 token。
9

5. 未來演進方向

在持續優化混合架構模型推理效能的進程中,SGLang正圍繞三大技術方向持續推進,嘗試拓展混合模型的應用邊界。
首先,在緩存機制的通用性與靈活性方面,社區已取得關鍵性突破。 升級後的MambaRadixTree現已全面支持Page Size > 1的靈活粒度配置,並實現了與MTP(Multi-Token Prediction)、Overlap Scheduler及Branching Position等先進機制的深度兼容。這一進展不僅有效解決了超長序列下的管理開銷問題,更顯著提升了內存利用效率,確立了系統對多樣化推理模式的高效適配能力。
在此堅實基礎上,阿里雲Tair KVCache將攜手SGLang重點推進HiCache分層KV緩存架構與混合模型的深度整合。 這不僅涉及多級混合存儲結構的重構,還需配套設計高效的存儲引擎查詢接口及緩存調度策略,旨在進一步提升緩存命中率,為混合模型在海量數據場景下提供低延遲、高吞吐的運行支撐。
最後,為保障模型在訓練與推理階段的嚴格一致性,團隊將持續推進比特級確定性推理的適配工作。 期望通過消除非確定性操作導致的數值偏差,進一步提升實驗的可復現性與生產部署的可靠性,完成從“高性能”到“高可信”的閉環。
參考鏈接:
[1]SGLang Hybrid Models:https://pytorch.org/blog/hybrid-models-meet-sglang-more-than-...

user avatar
0 位用戶收藏了這個故事!

發佈 評論

Some HTML is okay.