Hybrid Model Support:阿里云 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 的工作基础上深入剖析其设计原理、实现难点与优化路径,为基于混合架构的高效LLM 推理架构提供实践参考。

1.2 状态空间模型:线性效率与有限容量的权衡

状态空间模型(State Space Models, SSMs),通过递归式上下文压缩技术,将动态变化的 token 序列映射为固定维度的隐式状态。这种设计在计算范式上实现了双重突破:

1)内存效率提升:推理过程中状态维度恒定(),摆脱传统注意力机制随序列长度线性膨胀()的内存瓶颈;

2)计算复杂度降低:自回归生成时计算量仅随序列长度线性增长(),相较注意力机制的平方级复杂度()实现数量级优化。

然而,这种设计存在潜在约束:有限的状态容量与不可逆的压缩机制。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的回答就需要重头开始计算,前缀缓存失效。

在分布式部署层面,当前主流的 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,还实现了两类内存需求的正交管理,显著提升了整体内存利用率。

2.2 弹性内存池

然而,固定比例的池划分难以适应真实场景中波动的工作负载。例如,当系统负载从高并发的短对话任务切换至低并发但上下文极长的生成任务时,Mamba 池往往因请求减少而闲置,而 KV Cache 池却因长序列缓存需求激增而迅速耗尽,进而限制批处理规模引发性能瓶颈。

为此,SGLang 在多池隔离架构的基础上引入了弹性内存池机制,在保持 Mamba 状态与 KV 缓存语义隔离的前提下,实现池间容量的运行时动态重分配。该机制首先在存储管理层依托CUDA虚拟内存管理能力:系统在启动时预分配一个超额预定的虚拟地址空间,并为每个内存池创建可弹性伸缩的张量数据结构。这些张量本身不立即占用物理显存,而是作为虚拟占位符。当某类缓存需求增长时,控制模块将物理显存页动态映射到对应的虚拟地址区间,实现"按需激活"的内存分配;反之,当某内存池使用率下降,其空闲块所占的物理页会被解除映射并释放,从而回收资源。以长文本生成为例,当负载由短请求转为长序列任务时,推理批大小通常减小,SSM 层所需的 SSM 状态总量随之降低,Mamba 池使用率下降,系统便可自动将其空闲物理页转移至KV Cache池,支持更长上下文的持续扩展,有效缓解静态分配导致的内存利用率不均问题。

在控制决策层面,系统通过一个集中式调度模块实现智能、安全的池间资源再分配。各内存池在初始化阶段向该模块注册元信息。运行时,若某一池因容量不足发起扩容请求,控制模块会实时评估所有池的当前使用率,选择最空闲的池触发缩容操作------即释放其部分物理显存页,并在确认释放成功后,授权请求方完成扩容。整个过程严格限定在固定的总 GPU 显存预算内,无需重启服务或重新分配全局内存,既避免了 OOM 风险,又保障了分配操作的原子性与安全性。

通过"多池隔离 + 弹性调度"的协同设计,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状态则采用更灵活的弹性驱逐策略,允许从任意层级节点释放内存。

通过这一设计,MambaRadixCache 帮助混合了 SSM 和注意力层的模型能够在无需修改任何底层算子或 Mamba 推理算子的前提下无缝集成高效的前缀缓存能力。该方案不仅保持了原始计算逻辑的简洁与高性能,还显著降低了重复计算开销与内存占用,为大规模高效推理提供了关键基础设施支持。

3.2 推测解码适配方案

推测解码作为大模型推理加速的核心技术,在全注意力架构中通过并行生成并验证候选Token序列,显著提升了推理效率。然而,当将其应用于状态空间模型时,却面临根本性的适配挑战。其根源在于 SSM 的状态更新机制与传统注意力中的 KV Cache 存在本质差异:SSM采用原地更新策略,每处理一个新 Token,其内部状态更新可以简单抽象为递推公式:

,会被不可逆地覆盖。这种设计虽然在序列建模中高效简洁,却使得系统在推测解码的验证阶段无法像处理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 显式记录其在推测树中的父节点;进入验证阶段后,依据该索引追溯至对应的父状态,并执行递归更新。这一设计不仅保留了Eagle-Tree的多路径探索能力,还使其与SSM的状态演化机制对齐,成功将高效的推测解码扩展至SSM架构,为其实时推理提供了可行路径。

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标识符,并在后端传输模块中添加对应的读写、传输处理逻辑。整个过程高度模块化,无需侵入核心调度流程。

通过上述机制,SGLang实现了对异构模型状态的统一、高效且可扩展的管理,不仅兼容传统Transformer架构,也能无缝支持Mamba、SWA等新兴注意力变体,为混合架构大模型的高性能分离式推理提供了坚实基础。

4. 性能验证

SGLang 在v0.5.5版本用H200跑Qwen3-Next-80B-A3B-Instruct-FP8的实验验证了上述设计的有效性。如下图所示,启用前缀匹配以存代算的能力可以避免重复计算匹配的前缀,将TTFT降低至57.63%

对 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。

5. 未来演进方向

在持续优化混合架构模型推理效能的进程中,SGLang正围绕三大技术方向持续推进,尝试拓展混合模型的应用边界。

首先,在缓存机制的通用性与灵活性方面,社区已取得关键性突破。 升级后的MambaRadixTree现已全面支持Page Size > 1的灵活粒度配置,并实现了与MTP(Multi-Token Prediction)、Overlap Scheduler及Branching Position等先进机制的深度兼容。这一进展不仅有效解决了超长序列下的管理开销问题,更显著提升了内存利用效率,确立了系统对多样化推理模式的高效适配能力。

在此坚实基础上,阿里云Tair KVCache将携手SGLang重点推进HiCache分层KV缓存架构与混合模型的深度整合。 这不仅涉及多级混合存储结构的重构,还需配套设计高效的存储引擎查询接口及缓存调度策略,旨在进一步提升缓存命中率,为混合模型在海量数据场景下提供低延迟、高吞吐的运行支撑。

最后,为保障模型在训练与推理阶段的严格一致性,团队将持续推进比特级确定性推理的适配工作。 期望通过消除非确定性操作导致的数值偏差,进一步提升实验的可复现性与生产部署的可靠性,完成从"高性能"到"高可信"的闭环。

6. 了解更多

欢迎搜索钉钉群号:109765011301入群与技术专家交流!

相关推荐
Gobysec4 分钟前
Goby 漏洞安全通告|MindsDB /api/sql/query 未授权访问漏洞(CVE-2025-68472)
数据库·sql·安全
m0_748245925 分钟前
SQLite 数据类型概述
java·数据库·sqlite
五阿哥永琪7 分钟前
MySQL 回表查询 性能代价?如何避免?
数据库·mysql
DBA小马哥9 分钟前
文档型数据库MongoDB迁移替换至金仓数据库上线流程周期全解析
数据库·mongodb·文档型数据库
冰暮流星17 分钟前
sql语言之where语句
java·数据库·sql
橘子1321 分钟前
MySQL基础(一)
数据库·mysql·php
難釋懷30 分钟前
安装Redis
数据库·redis·缓存
jiayong2332 分钟前
Word协作与审阅实用手册
服务器·数据库·word
涵涵(互关)33 分钟前
添加了 @TableId(type = IdType.AUTO) 但仍生成超大 ID
数据库·spring·mybatis
什么都不会的Tristan1 小时前
redis-原理篇-SDS
数据库·redis·缓存