SGLang颠覆！大模型内存O(1)，推理狂飙324Token/秒！

在当下人工智能浪潮汹涌而至的时代，大型语言模型（LLM）正以前所未有的速度改变着我们的生活和工作方式。然而，随着模型规模的不断扩大和上下文长度的日益增长，如何在实际应用中高效、经济地部署和运行这些模型，成为了一个摆在所有开发者面前的巨大挑战。特别是对于那些需要处理长文本、执行复杂推理任务（如RAG查询、智能体工具链）的场景，模型的服务性能和内存消耗常常成为瓶颈。

传统的全注意力层模型在处理长上下文时，其计算复杂度和内存消耗都会呈二次方增长，这无疑给服务器资源带来了巨大的压力。为了突破这一瓶颈，业界开始探索混合模型架构，将全注意力层的强大能力与Mamba、线性注意力等高效替代方案结合起来。这些混合模型在保证性能的同时，极大地优化了资源占用。今天，就让我们一同深入了解SGLang，看看它如何创新性地应对这些挑战，为大模型的高效服务开辟新路径。新媒网跨境获悉，SGLang在混合模型优化方面取得了显著进展。

揭秘状态空间模型（SSM）：大模型效率的秘密武器

要理解SGLang的优化，我们首先要认识状态空间模型（SSM）。SSM，以及更广义的线性循环神经网络（RNN）和线性注意力机制，例如Mamba模型，它们的核心优势在于能够选择性地将文本中的信息和上下文压缩成一个固定大小的循环状态。这个循环状态会实时更新，并且其大小是恒定的。

相比之下，传统的全注意力机制随着输入序列长度的增加，其计算复杂度和内存消耗都会以平方级别增长，这在处理超长文本时效率低下。而采用SSM后，内存消耗可以保持在恒定水平，计算复杂度则与序列长度呈线性关系。这无疑为长上下文模型的部署带来了革命性的变化。

不过，纯粹的线性结构虽然高效，但也存在局限性，其有限的状态容量在处理极长上下文或实现强大的召回能力时可能会遇到挑战。为了在效率和能力之间找到最佳平衡点，混合模型应运而生。这类模型巧妙地将二次方复杂度的注意力层与SSM层交错使用，在保持SSM层大部分效率优势的同时，确保了模型在各种任务上都能表现出色。

我们可以通过一个简单的对比，直观地看到SSM在计算复杂度和内存消耗上的优势：

特性	注意力模型 (Attention)	状态空间模型 (SSM)
计算复杂度	O(N²)	O(N)
内存消耗	O(N)	O(1)

混合状态管理的精妙设计

在SGLang的架构中，针对混合线性模型，内存池被巧妙地划分为两个独立的部分：Mamba内存池和KV缓存内存池。这种设计的一个核心优势在于，两个内存池的大小都是预先固定好的，从而从根本上消除了CUDA内存溢出（OOM）的风险，确保了系统运行的稳定性。

更值得一提的是，用户可以根据实际的工作负载特性，通过调整服务器参数--mamba-full-memory-ratio来灵活控制Mamba内存池和KV缓存内存池之间的比例。这种可配置性赋予了系统极大的灵活性，能够更好地适应不同应用场景的需求。

Mamba内存池与KV缓存内存池的主要区别在于，前者是按请求级别分配Mamba状态的，而后者则是按令牌（token）级别进行分配。为了有效地将Mamba状态与具体的请求绑定，SGLang引入了HybridReqToTokenPool机制，确保请求的生命周期与Mamba状态的生命周期保持一致。

此外，通过HybridLinearKVPool，SGLang能够将逻辑层ID映射到KV缓存内存池中的实际层索引，这意味着在线性层中无需为KV缓存进行额外的分配，从而大幅节省了内存空间。这一系列精巧的内存管理设计，共同构筑了SGLang高效稳定运行的基础。

弹性内存池：智慧调配，释放硬件潜力

混合模型集成了多种注意力类型，每种类型都有自己独立的内存池，并配备专门的分配器进行GPU内存管理。为了最大化内存利用率并优化推理性能，必须根据工作负载的特点来配置这些内存池之间的比例。然而，手动设置这些比例并非易事，而且工作负载的动态变化可能导致预设比例在运行时变得不那么理想。

为了解决这一难题，SGLang提出了一个创新的弹性内存池方案。在固定的GPU总内存预算下，这个弹性内存池能够动态调整池的大小。它主要由两部分组成：

基于CUDA虚拟内存管理的可伸缩张量： SGLang预先分配了一个具有超额容量的虚拟地址空间。然后，在这个空间内创建一个torch.Tensor，并将其重塑以匹配KV缓存的需求。当需要扩展内存池时，物理的CUDA内存页会被映射到相应的虚拟地址，从而激活对应的KV缓存块。反之，当需要收缩内存池时，空闲的KV缓存块会被禁用，它们的物理页也会被解除映射以释放内存。
集中控制模块： 在系统初始化时，所有的内存池都会向这个控制模块注册。在运行时，如果某个内存池的容量耗尽，它会向控制模块请求扩容。控制模块会识别出当前利用率最低的内存池，发出收缩指令，并在成功收缩后授权请求者进行扩容。

通过引入弹性内存池，SGLang系统能够根据工作负载的需求，动态调整Mamba内存池和KV缓存内存池之间的分配比例，最大限度地提高GPU内存利用率，从而支持更大批量的推理任务。这种智慧的内存调配机制，为大模型在复杂多变的应用场景中持续高效运行提供了坚实保障。

关键优化与适配：让大模型更聪明高效

在追求大模型极致性能的道路上，SGLang并未止步于内存管理，更在关键优化技术上进行了深度探索与创新。

1. 前缀缓存：加速重复请求

前缀缓存（Prefix Caching）是全注意力模型中一项被广泛使用的优化技术，它能够有效避免跨请求的冗余计算，显著提升服务效率。想象一下，如果多个用户的查询都以相同的长前缀开始，前缀缓存就能让模型避免重复计算这部分前缀，直接从预存的状态继续处理。

然而，对于Mamba状态来说，实现前缀缓存却面临一些特殊挑战：

状态原地更新： Mamba状态是原地更新的，这意味着一旦更新，之前的状态就无法回溯，也就难以直接表示其前缀。
状态规模庞大： Mamba状态的大小通常比单个令牌的KV缓存大一个数量级，这增加了存储和管理的复杂性。
"全有或全无"的复用特性： 大多数Mamba状态的前向核函数在复用时表现出“全有或全无”的特性，这进一步限制了其在通用前缀缓存中的应用。

面对这些挑战，SGLang通过引入一种名为MambaRadixCache的混合基数树结构，成功地为混合线性模型实现了前缀缓存。MambaRadixCache主要包括以下三个核心部分：

匹配（match）： 当有新的输入请求时，MambaRadixCache会返回最佳匹配节点。如果该节点的Mamba状态值不为空，并且其键是输入的前缀，系统就会从基数树中复制Mamba状态。
插入（insert）： 在分块预填充（chunked prefill）或解码阶段完成后，KV缓存和Mamba状态会插入到MambaRadixCache中。这需要从请求中派生出Mamba状态的一个检查点（checkpoint）。
淘汰（evict）： MambaRadixCache维护了两个独立的LRU（最近最少使用）列表，分别用于管理Mamba状态和KV缓存的时间戳。KV缓存必须从叶节点向根节点逐级淘汰，而Mamba状态则可以从任何节点进行淘汰。

通过集成MambaRadixCache，混合线性模型无需修改线性注意力核函数，就能享受到前缀缓存带来的巨大性能提升。

2. 推测解码：提速模型生成

推测解码（Speculative Decoding）是一种通过预测未来令牌来加速模型生成过程的有效方法。但在SSM模型中，这种技术同样遇到了难题。

状态原地更新： 如前所述，SSM状态的原地更新特性使得被拒绝的令牌无法回滚，这是标准推测解码面临的主要障碍。
与Eagle-Tree注意力掩码不兼容： Eagle-Tree的注意力掩码机制与SSM状态的维护方式存在不兼容性。

SGLang的解决方案巧妙地为每个草稿令牌（draft token）分配了一个独立的Mamba缓存槽。这意味着每个草稿令牌都有自己独立的SSM状态。例如，当模型生成“the air streets”时，“the”会有一个槽位，“air”会有一个槽位，“streets”也会有一个槽位。当一个序列的草稿令牌被接受后，系统只需简单地将最后一个被接受的槽位提升为新的主状态。举例来说，如果“the streets are”被接受，那么“streets”对应的槽位就会成为新的主SSM状态。

结合EAGLE-Tree，并采用Top-K > 1的策略，系统在验证之前会预先计算父索引。对于每个草稿令牌，它通过这些索引追溯其父级，然后应用递归更新Snew = Sparent + vnew knewᵀ。新媒网跨境了解到，这种创新方法极大地提升了推测解码在SSM模型中的效率。

3. 预填充与解码分离：优化传输效率

SGLang的预填充（Prefill）与解码（Decode）分离（PD disaggregation）架构，通过扩展传输协议，引入了一个专用的状态传输通道，从而完美支持混合模型。除了标准的页式KV缓存传输之外，系统还通过并行数据路径传输模型特定的状态（例如Mamba的卷积/时间状态、SWA滑动窗口）。

具体到Mamba模型的集成细节：

Mamba模型维护两个独立的内存池：一个用于全注意力层的页式KV池，另一个用于线性层，存储卷积和时间状态的Mamba池。
当新请求到达时，它首先通过MambaRadixTree进行前缀匹配。如果命中缓存，匹配到的MambaState会被复制到一个新的Mamba内存区域，作为当前请求的Mamba缓冲区，预填充推理将在此基础上继续进行。
预填充完成后，预填充实例会将最终的Mamba状态作为一个单一的连续块传输给解码实例，使用dst_state_indices来标识目标槽位。与可以增量流式传输的页式KV传输不同，Mamba状态是原子性传输的。
解码实例会预先分配KV页面槽位和专用的Mamba槽位，确保接收到的状态存储在正确的内存位置，以便后续的解码步骤使用。

要将一个新的混合池集成到现有的PD分离服务中，只需三个简单步骤：

暴露状态缓冲区指针、大小和项目长度，以进行传输注册。
在预填充和解码工作器中定义state_indices准备逻辑，以指定要传输哪些池槽位——这可以是一个请求的单个索引（例如Mamba），窗口化数据的页面索引（例如SWA），或完整的序列索引（例如NSA）。
在KV管理器中注册一个唯一的state_type标识符，并在后端添加相应的传输处理。