引言
混合模型结合了全注意力层与替代方案（如 Mamba 或线性注意力）的优势，在长上下文大语言模型（LLM）推理场景中越来越受欢迎。通过利用线性注意力，每个请求的 KV 缓存内存消耗被限制在一个常数内，且预填充（prefill）延迟可以随输入长度线性扩展。这一特性非常契合实际生产环境中的工作负载，例如 RAG 查询、智能体工具调用以及思考/推理模式。
然而，原地状态更新（in-place state updates）使得无法针对部分序列匹配回滚缓存条目，这增加了许多广泛采用的功能的实现难度，例如前缀缓存（prefix caching）和推测解码（speculative decoding）。Mamba 状态所需的请求级状态存储也为内存管理和 PD（预填充-解码）存算分离带来了新的挑战和需求。
本文将讨论 SGLang 如何适应并优化上述挑战。
什么是状态空间模型？
状态空间模型（SSM）以及更广泛的线性 RNN 和线性注意力（如 Mamba），通过将 Token 和上下文选择性地压缩到循环状态中来实现功能。该循环状态大小固定且原地更新。通过使用 SSM，内存消耗可以保持在恒定水平，计算复杂度随序列长度线性扩展，而非平方级。然而，纯线性结构受限于有限的状态容量，在处理超长上下文或实现强大的召回能力方面存在挑战。
为了在效率和容量之间取得平衡，混合模型应运而生。这些模型以固定间隔将平方级的注意力层与 SSM 层交替排布。因此，混合模型在各种任务中都能实现卓越的性能，同时保留了 SSM 层所提供的大部分效率优势。
| 注意力 | SSM | |
| 计算复杂度 | O(N^2) | O(N) |
| 内存消耗 | O(N) | O(1) |
内存管理
混合状态管理设计
在 SGLang 中，混合线性模型将内存池分为两部分：Mamba 池和 KV 缓存池。Mamba 池和 KV 缓存池的内存大小均为固定值，从而消除了 CUDA 显存不足（OOM）错误的风险。用户可以通过更改服务器参数 –mamba-full-memory-ratio，根据工作负载调整 Mamba 池和 KV 缓存池之间的比例。
Mamba 池和 KV 缓存池的主要区别在于，前者在请求级别分配 Mamba 状态，而后者在 Token 级别进行分配。我们使用 HybridReqToTokenPool 将 Mamba 状态与请求绑定，使请求与 Mamba 状态的生命周期保持一致。此外，我们使用 HybridLinearKVPool 将逻辑层 ID 映射到 KV 缓存池中的实际层索引，这样我们无需在线性层中分配 KV 缓存，从而大幅节省了内存空间。
弹性内存池
混合模型集成了多种注意力类型，每种类型都有自己的内存池，并配备了专门的分配器进行 GPU 内存管理。为了最大化内存利用率并优化推理性能，必须根据工作负载特性配置这些内存池之间的比例。然而，手动设置这些比例并非易事，且波动的负载可能会使预设比例在运行时变得不再最优。为了解决这个问题，我们提出了一种弹性内存池，它能在固定的总 GPU 内存预算下动态调整池大小。
弹性内存池由可调整大小的张量和中央控制模块组成
通过 CUDA 虚拟内存管理实现可调整大小的张量（Resizable Tensors）：
- 预先分配具有超额容量的虚拟地址空间。在此空间内创建一个 torch.Tensor，并根据 KV 缓存需求进行重塑（reshape）。
- 若要扩展内存池，物理 CUDA 内存页面将被映射到相应的虚拟地址，从而激活对应的 KV 缓存块。
- 若要收缩内存池，空闲的 KV 缓存块将被禁用，其物理页面将被取消映射以释放内存。
中央控制模块
- 在初始化期间，所有内存池都会在控制模块中进行注册。
- 运行时，如果某个内存池耗尽了容量，它会请求扩展。控制模块将识别出利用率最低的池，发出收缩指令，并在成功收缩后授权请求方进行扩展。
通过弹性内存池，系统能够根据工作负载需求动态调整 Mamba 池和 KV 缓存池之间的分配比例，从而最大限度地提高 GPU 内存利用率，实现更大批量的推理。
优化与适配
前缀缓存（Prefix Caching）
前缀缓存是全注意力模型中广泛使用的一种优化方法，可以节省跨请求的重复计算。然而，Mamba 状态的以下特性使得前缀缓存变得复杂：1) SSM 状态是原地更新的，因此请求的状态无法回滚以表示其前缀；2) SSM 状态比单个 Token 的 KV 大几个数量级；3) 大多数 SSM 状态的前向内核表现出“全有或全无”的可重用性。
SGLang 通过实现名为 MambaRadixCache 的混合基数树来支持混合线性模型的前缀缓存。它主要分离了匹配/插入/驱逐部分：
- 匹配 (match)：MambaRadixCache 将返回最佳节点，其中 Mamba 状态值不为 None 且 Key 为输入的前缀。它需要从基数树中复制 Mamba 状态。
- 插入 (insert)：在分块预填充或解码阶段之后，KV 缓存和 Mamba 状态将被插入到 MambaRadixCache 中。这需要从请求中分叉（fork）出 Mamba 状态的检查点。
- 驱逐 (evict)：MambaRadixCache 保留两个 LRU 列表，分别维护 Mamba 状态和 KV 缓存的时间戳。KV 缓存必须从叶子节点到根节点进行驱逐，而 Mamba 状态可以从任何节点驱逐。
通过集成 MambaRadixCache，混合线性模型无需修改线性注意力内核即可使用前缀缓存。
推测解码（Speculative Decoding）
为简单起见，我们使用最基础的线性注意力更新来阐述：
Sₜ = Sₜ₋₁ + vₜ kₜᵀ，
以便于阅读。在实际系统中，更新过程会稍微复杂一些。
为什么标准的推测解码不适用于 SSM？
- SSM 状态原地更新，因此被拒绝的 Token 无法回滚。
- Eagle-Tree 注意力掩码与 SSM 状态的维护方式不兼容。
SGLang 的解决方案：每个草稿 Token 使用一个独立的 Mamba 缓存槽
- 每个草稿 Token 都会获得一个私有的缓存槽，其中包含其自己的 SSM 状态
- “the” → 槽 1
- “air” → 槽 2
- “streets” → 槽 3
- 当草稿 Token 序列被接受时，只需将最后一个接受的槽提升为新的主状态。
（示例：在接受 “the streets are” 后，槽 3 成为主 SSM 状态。）
使用 Top-K > 1 的 EAGLE-Tree
- 在验证前预计算父索引。
- 对于每个草稿 Token
- 使用这些索引追踪其父节点。
- 应用递归更新 Snew = Sparent + vnew knewᵀ
预填充与解码存算分离（Prefill and Decode disaggregation）
SGLang 的 PD-存算分离架构通过增加专门的状态传输通道来支持混合模型。除了标准的 Paged KV 缓存传输外，系统还通过并行数据路径传输模型特定的状态（例如 Mamba 卷积/时间状态、SWA 滑动窗口）。
Mamba 集成细节：
- Mamba 模型维护两个独立的内存池：用于全注意力层的分页 KV 池，以及用于存储卷积和时间状态的线性层 Mamba 池。
- 当新请求到达时，它首先通过 MambaRadixTree 进行前缀匹配。如果缓存命中，匹配的 MambaState 将被复制到新的 Mamba 内存区域，作为当前请求的 Mamba 缓冲区，预填充推理随即继续进行。预填充完成后，预填充实例将最终的 Mamba 状态作为一个连续块传输给解码实例，并使用 ‘dst_state_indices’ 标识目标槽。与可以增量流式传输的 Paged KV 传输不同，Mamba 状态是原子传输的。
- 解码实例预分配 KV 页面槽和专门的 Mamba 槽，确保接收到的状态存储在正确的内存位置，以供后续解码步骤使用。
要在当前 PD 实现中集成新的混合池，只需三个步骤：
- 暴露状态缓冲区指针、大小和条目长度以进行传输注册；
- 在预填充和解码 worker 中定义 state_indices 准备逻辑，以指定传输哪些池槽——这可以是每个请求的单个索引（例如 Mamba）、窗口数据的页面索引（例如 SWA），或完整序列索引（例如 NSA）；
- 在 KV 管理器中注册一个唯一的 state_type 标识符，并在后端添加相应的传输处理逻辑。
基准测试
基准测试是在 SGLang v0.5.5（最新发布版本）上进行的。服务器运行在 H200 GPU 上，模型为 Qwen3-Next-80B-A3B-Instruct-FP8。
前缀缓存（Prefix Caching）
python3 -m sglang.launch_server --model Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 –tp 2 python3 -m sglang.bench_serving --backend sglang \ --dataset-name generated-shared-prefix \ --gsp-num-groups 50 \ --gsp-prompts-per-group 10 \ --gsp-system-prompt-len 10240 \ --gsp-question-len 256 \ --gsp-output-len 128 \ --max-concurrency 5 --port 30000
推测解码（Speculative Decoding）
python3 -m sglang.launch_server –model Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 –tp 2 –disable-radix-cache –speculative-num-steps 2 –speculative-eagle-topk 1 –speculative-num-draft-tokens 3 –speculative-algo EAGLE
python3 -m sglang.test.send_one
我们在 Batch size = 1 的情况下测试了 Qwen3-Next-80B-A3B-Instruct-FP8 的性能。在 2-token MTP 窗口且 topk=1 时，系统达到了 257.20 tokens/sec 的吞吐量，平均接受长度为 2.709 个 Token。
在 3-token MTP 窗口且 topk=1 时，吞吐量提高至 306.94 tokens/sec，平均接受长度为 3.413 个 Token。
在 4-token MTP 窗口、topk=4 且 draft tokens=8 时，吞吐量提高至 324.57 tokens/sec，平均接受长度为 4.231 个 Token。
未来工作
后续工作可在 此处 跟踪。具体而言，我们计划：
- 更通用的前缀缓存：包括支持 page size > 1、推测解码及其他功能。
- 集成到 Hicache：快速分层 KV 缓存是 SGLang 的重要功能。我们需要在线性注意力层中为 KV 缓存开发新的查询、存储和调度机制。
- 确定性推理适配：我们希望对混合线性模型进行适配，以支持位级训练-推理一致性（bitwise training-inference consistency）。
致谢
SGLang 团队：Yi Zhang, Biao He, Binyao Jiang, Ke Bao, Qingquan Song, Hanming Lu, Shangming Cai, Zhangheng Huang, Sicheng Pan, Baizhou Zhang