Disaggregated Serving for Hybrid SSM Models in vLLM

发表时间: 2026-04 · Blog post by Nicolò Lucchesi, Zhanqiu Hu (Red Hat), and the vLLM team (blog.vllm.ai)

原文: https://blog.vllm.ai/blog/2026-04-21-hybrid-ssm-disagg

作者/机构：Nicolò Lucchesi, Zhanqiu Hu (Red Hat), and the vLLM team

速读

一句话结论 本文通过扩展 vLLM 的网络传输组件，提出了一套双视图描述符和连续内存布局方案，成功在交替包含注意力层和状态空间层的混合模型上实现了分离式预填充与解码架构，并在 120B 参数模型的高并发场景下实现了吞吐量对同置部署基线的全面超越。

要解决什么问题 原有分离式推理架构（通过 RDMA 网络直接传输键值对缓存）卡在无法支持交替包含全注意力（FA）层和状态空间（SSM）层的混合模型。这个卡点至关重要，因为 FA 层和 SSM 层需要保存的上下文状态在机制上存在根本差异：FA 层保存的是随 token 数量增长的键值对缓存，而 SSM 层（如 Mamba）保存的是固定大小的序列历史摘要（包含卷积状态和时间状态）。为了节省显存，vLLM 的混合内存分配器（HMA）会将这两种层的状态塞进同一个物理张量池中，并通过填充无用字节来强行对齐它们的内存页大小。然而，底层的网络传输组件只认单一且均匀的缓存格式，无法用同一套“地址加长度”的规则去同时读取变长的 FA 块和定长的 SSM 块。此外，在预填充和解码实例的张量并行度不一致时，SSM 的卷积状态在内存中是交错排列的，如果直接通过网络拉取，要么需要发起大量非连续的低效内存读取，要么必须在接收端分配与完整状态等大的临时显存来做数据重排，这会严重挤占本应用于生成 token 的显存空间。

怎么做的 核心思路是解耦物理内存与网络传输的逻辑视图，并通过重构张量布局实现零拷贝的切片传输。这套方案完全绕开了统一格式的限制和显存重排的开销，主要由三个关键部件构成。第一是双描述符视图，它在同一块物理内存上注册两套独立的网络传输描述符列表，前半部分按 FA 层的真实长度和偏移进行索引，后半部分按 SSM 层的真实长度进行索引。当解码端需要拉取数据时，根据层类型加上对应偏移量即可准确命中数据，其映射逻辑为：

$$\begin{aligned} desc\_id = \begin{cases} region\_id \times N_{phys} + block\_id, & \text{if FA group} \\ mamba\_region\_id \times N_{log} + block\_id + num\_descs, & \text{if Mamba group} \end{cases} \end{aligned}$$


第二是物理与逻辑块桥接，由于注意力算子强制要求特定的物理块大小（如 16 个 token），而 SSM 没有 token 维度，方法对 FA 层应用物理块拆分比例，而对 SSM 层保持原始逻辑块计数，两者在同一个传输句柄下独立寻址。第三是针对异构张量并行的三描述符卷积传输，这是解决显存开销的核心。方法强制将 SSM 的卷积状态切换为 DS 布局，即维度在前、状态长度在后：(dim, state_len)。在这种布局下，卷积状态内部的三个子投影（x、B、C）在内存中变成了连续的块。当解码端张量并行度大于预填充端时，每个解码节点只需计算自己负责的偏移量，用三个连续的读取指令直接把属于自己的切片拉取到最终的显存位置。这种设计既不需要在接收端开辟巨大的临时缓冲池，也不需要执行任何显存重排算子，同时还完美跳过了底层内存分配器为了对齐页大小而塞入的填充字节，实现了真正按需传输的零开销。

效果如何 实验在 8 张通过 NVLink 连接的 H200 显卡上搭建，测试模型为 120B 参数规模的混合架构模型 nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8，评测数据使用 ShareGPT 数据集。对比基线是同置部署路线（Co-located baseline），即单个实例独占 8 张卡运行张量并行度为 8 的完整推理；本文方法代表分离式路线，将 8 张卡拆分为 1 个张量并行度为 4 的预填充实例和 1 个张量并行度为 4 的解码实例。量化结果表明，在 8 到 256 个并发用户的吞吐量与延迟测试中，本文的分离式架构在较高并发下 Pareto 占优于同置部署基线，通过隔离预填充对解码的干扰，解码实例能维持更大的批处理大小，显著提升了单卡的输出吞吐量。在传输代价方面，得益于跳过填充字节的优化，在 bf16 精度设置下，该方法每条请求能消除约 50MB 的无效网络传输。作者也承认了该方法的局限性：目前仅支持 Mamba2 模型（因为 Mamba1 的状态形状无法重建所需的中间维度），尚未充分验证与投机解码结合时的稳定性，且在开启混合内存分配器时，暂不支持预填充和解码两端使用不同的逻辑块大小。

主要贡献

当前，将Mamba风格的状态空间模型（SSM）层与标准的全注意力（FA）层交替使用的混合架构（例如NVIDIA Nemotron-H）正受到广泛关注，这种架构结合了状态空间模型在线性时间上的效率以及注意力机制的表达能力。vLLM系统此前已经通过基于NIXL的KV连接器支持了标准Transformer模型的分离式Prefill/Decode（P/D）服务：即Prefill实例计算KV缓存块，Decode实例通过RDMA拉取这些块，从而消除冗余的重复计算。然而，将这一特性扩展到混合模型面临着核心挑战：FA层和SSM层存储的状态在本质上是不同的，它们具有不同的布局和大小，而现有的块管理器和NIXL连接器是围绕单一、统一的KV缓存格式设计的。

为了解决这一问题，本文扩展了NIXL连接器，使其能够在分离模式下支持混合SSM-FA模型。本文的核心创新点包括：

  1. 双描述符视图（Dual descriptor views）：设计了两套NIXL块描述符，它们以不同的偏移量和大小索引相同的物理内存区域，一套用于FA块，另一套用于SSM块。
  2. 物理与逻辑块桥接（Physical/logical block bridging）：处理了块管理器所见的逻辑块抽象与注意力算子所需的物理块大小之间的不匹配问题。
  3. 3-描述符卷积传输（3-descriptor conv transfer）：对Mamba卷积状态进行分解，使得在异构张量并行（Heterogeneous-TP）传输时，无需在发送端重新洗牌（reshuffling）数据。

这些更改完全是附加的扩展功能，不会修改标准Transformer模型的现有工作流，仅在模型包含SSM层时激活。该功能在vllm>=v0.20.0版本中可用。

背景知识与设计挑战

NIXL KV传输工作流：在标准Transformer中，NIXL分离式P/D的工作流包含四个阶段。首先是注册内存区域，每个工作节点向NIXL注册其KV缓存张量以便通过RDMA访问。其次是创建块描述符，为每个注册区域创建指定$(address, length, device\_id)$的逐块描述符，这是传输的基本单位。接着是握手阶段，当Decode（D）工作节点首次需要从Prefill（P）工作节点拉取数据时，双方交换代理句柄、块数量、块长度等元数据。最后是传输阶段，调度器指示D从P拉取哪些块，D将$block\_id$映射为$descriptor\_id$，发出RDMA READ请求并轮询完成状态。对于具有$M$个注册区域和$N$个块的标准模型，区域$r$中的块ID $b$映射到描述符索引$r \times N + b$。混合模型的挑战在于这种统一的方案不再适用，因为FA层和SSM层需要不同的描述符大小和块数量。

FA与SSM状态的根本差异：在标准Transformer中，每一层的KV缓存形状相同，如$[num\_blocks, 2, block\_size, num\_kv\_heads, head\_dim]$，所有层共享相同的块大小、页面大小和块数。然而，Mamba层存储的内容截然不同，它不维护每个Token的K/V对，而是维护一个折叠的卷积状态（Conv state）和时间SSM状态（SSM state）。例如，卷积状态的形状为$(conv\_dim, state\_len)$，SSM状态的形状为$(num\_heads, head\_dim, state\_size)$。这些状态中没有“Token”的概念，它们是整个序列历史的固定大小摘要。这意味着SSM的$block\_size$实际上为1：每个块是一个完整的状态快照，而不是每个Token向量的组合。在此场景下，块（block）依然是单一的传输单位。

HMA共享张量布局机制：vLLM的混合内存分配器（HMA）按类型对层进行分组（所有FA层为一组，所有SSM层为另一组），然后跨组池化内存，使得各组中相同位置的层共享相同的物理张量。这种设计提高了效率，但也意味着同一个张量被一组视为FA块的同时，被另一组视为SSM块。由于FA页面的大小由$block\_size \times num\_kv\_heads \times head\_dim \times 2$决定，而SSM页面的大小为$conv\_state\_bytes + ssm\_state\_bytes$，HMA会增加FA的$block\_size$直到其大于Mamba的页面大小，然后对Mamba行进行填充（$+[pad]$），使得两组在字节上具有相等的页面大小。这给NIXL带来了问题：单个具有统一$(address, length)$条目的描述符列表无法正确索引这两种视图。FA描述符指向$base + b \times page\_size$且长度为$fa\_block\_len$，而Mamba描述符指向相同的基地址但长度为$conv\_size$或$ssm\_size$，两者存在显著差异，因此需要将它们注册在不同的描述符上，以允许在异构设置（即D的TP不等于P的TP）下索引KV头。

方法细节

双描述符视图设计：为了解决共享内存的索引问题，本方案在同一物理内存上注册了两个独立的描述符列表。这两个列表被拼接在一起，并由单个NIXL传输句柄指向。FA描述符占据了前$num\_descs = M \times N\_phys$个槽位（包含对K和V的分别索引），而Mamba描述符紧随其后（包含x、B、C和SSM的索引）。在映射块ID时，系统会进行判断：如果是FA组，则$desc\_id = region\_id \times N\_phys + block\_id$；如果是Mamba组，则$desc\_id = mamba\_region\_id \times N\_log + block\_id + num\_descs$。

物理与逻辑块大小的差异处理：注意力内核（如FlashInfer）通常需要特定的物理块大小（例如16个Token），这可能与用户设置或HMA计算的逻辑块大小不同。对于标准模型，这通过一个简单的比率处理：$physical\_blocks = logical\_blocks \times ratio$，其中$ratio = logical\_block\_size / kernel\_block\_size$。在混合模型中，这个比率仅适用于FA层。由于SSM层没有可以分割的“Token”维度，它们始终直接使用$logical\_blocks$。因此，描述符列表的FA部分和Mamba部分使用不同的块数：FA部分为$M \times N\_phys$个块（其中$N\_phys = N\_logical \times ratio$），而Mamba部分为$M \times N\_logical$个块。这一差异通过$\_physical\_blocks\_per\_logical$字段在每个引擎级别进行跟踪，并在$\_get\_block\_descs\_ids$的ID映射中使用适当的步幅。

3-描述符卷积传输方案：在同构张量并行（TP）下，传输SSM状态很简单；但在异构TP（例如$P\_TP=1, D\_TP=4$）下，四个D工作节点必须从单个P工作节点读取其对应的卷积和SSM状态分片。SSM时间状态沿头（heads）维度分片，切片很容易。但卷积状态结构为$[x | B | C]$（分别为子投影），在标准的SD布局$(state\_len, dim)$中，这些子投影在内存中是交错的。如果D工作节点只想要其$x$的部分，将需要收集不连续的字节，这对于零拷贝RDMA是不切实际的。

DS布局解决方案：为了解决内存交错问题，系统要求卷积状态使用DS布局$(dim, state\_len)$（通过设置环境变量$VLLM\_SSM\_CONV\_STATE\_LAYOUT=DS$实现）。在这种布局中，每个子投影的数据在单个页面内是连续的（即$|--- x ---|--- B ---|--- C ---|--- SSM ---|$）。因此，每个D节点可以通过三次独立的、连续的RDMA读取来获取其$x$、$B$和$C$的切片（底层仍仅发出一个NIXL READ请求）。对于异构TP，$remote\_conv\_offsets$方法计算每个D节点的切片在P页面中的位置，从而为每个Mamba层生成4个描述符区域（x, B, C, SSM）。虽然这增加了描述符列表的长度，但RDMA传输本身保持了高效的连续读取，且在GPU上不需要分配额外的内存暂存缓冲区，也无需在任何一侧重新洗牌数据。

零开销传输优化：本方案刻意避免了将整个卷积状态传输给每个D节点然后再在本地进行排列/切片的简单替代方案。首先，无需暂存缓冲区，避免了在每个D节点上分配与P的完整卷积状态大小相同的临时缓冲区（在Nemotron-H中，bf16格式下每个块的卷积状态达$3 \times 3072 \times 2$字节），从而节省了可用于KV缓存的宝贵空间。其次，传输后无需重新洗牌，D节点直接将所需字节读取到KV缓存的最终目的地。第三，按需传输，每个D节点仅传输其$1/TP$份额的卷积状态，对于$D\_TP=4$，数据量减少了4倍。最后，跳过HMA填充，Mamba描述符的尺寸被精确设置为实际的$conv\_bytes + ssm\_bytes$，这意味着永远不会通过网络传输HMA为匹配FA页面大小而添加的填充字节。图1验证了这一优化在Nemotron Super 120B模型上的效果，在bf16设置下，该方法每个请求消除了约50 MB的不必要传输。

Figure 1: P→D transfer volume vs. input sequence length for Nemotron Super 120B (TP=4, FA block_size=4224). The Naive and Optimal baselines are computed analytically from the model's page sizes and block counts. The Measured line reports the actual bytes transferred (as reported by NIXL) during disaggregated P/D serving. Our approach (Optimal) eliminates HMA padding overhead, which is reflected in the measured transfer.
图1：Nemotron Super 120B（TP=4，FA block_size=4224）的P→D传输量与输入序列长度的关系。Naive和Optimal基线是根据模型的页面大小和块数解析计算的。Measured线报告了在分离式P/D服务期间实际传输的字节数（由NIXL报告）。我们的方法（Optimal）消除了HMA填充开销，这反映在测量的传输中。

Nemotron-H 实例整合：以在TP=2下提供nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8的分离式P/D服务为例。该模型共有52层，Mamba与FA交替。HMA将它们分为5组（4个Mamba组，1个FA组），池化后产生6个共享KV缓存张量。在注册阶段，这6个共享张量被注册为NIXL内存区域，首先为所有6个区域创建FA描述符，随后附加带有4个子区域的Mamba描述符。在传输流中，P完成Prefill后，调度器分配块ID；D接收块ID并将其映射到描述符索引；随后D发出单个包含FA和Mamba描述符的make_prepped_xfer READ请求，并轮询完成。从D的视角看，整个传输是一个单一的异步操作，没有中间缓冲区，也没有数据重新洗牌。

实验环境

  • 硬件配置：8张通过NVLink连接的H200 GPU。
  • 模型架构：nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8。这是一个近期的120B LatentMoE混合架构模型，交替包含Mamba2和全注意力（FA）层。
  • 数据集：ShareGPT（用作测试数据集）。

  • 对比配置：

    • 同位服务基线（Co-located baseline）：单一实例，TP=8，使用全部8张GPU。
    • 分离式P/D（Disaggregated P/D）：1个Prefill实例（TP=4，使用4张GPU）+ 1个Decode实例（TP=4，使用4张GPU），总GPU数量相同。

  • 其他软件/运行设置：关闭前缀缓存（Prefix-caching disabled）。所有运行均使用非常高的预热（warmup）值，以确保KV缓存被“打乱（scrambled）”，从而避免当请求块碰巧连续分配时获得的初始性能提升，这能更准确地反映常规的长时间运行使用情况。

实验结果

  • 实验内容：在并发用户数从8扫描到256的范围内，测试并绘制了每GPU输出吞吐量（Output throughput per GPU）与单用户输出Token速率（Interactivity）之间的帕累托曲线（Pareto curve）。
  • 实验结果与分析：结果表明，在混合SSM模型中，分离式P/D在较高批处理大小下在帕累托曲线上主导了同位服务基线（见图2）。通过将Decode阶段与Prefill阶段的干扰隔离开来，Decode实例可以在不发生停滞的情况下维持更大的批处理，从而在高并发下显著提高每GPU的输出Token/秒。这与在标准Transformer模型中观察到的分离式服务模式一致。

Figure 2: Disaggregated P/D vs. co-located serving for a hybrid SSM model. Throughput-vs-latency Pareto curve across concurrency levels. Prefix-caching disabled.
图2：混合SSM模型的分离式P/D与同位服务的对比。跨并发级别的吞吐量与延迟帕累托曲线。前缀缓存已禁用。

结论

本文成功地将vLLM中的NIXL KV连接器扩展到了混合SSM-FA模型的分离式服务中，解决了FA与SSM状态在内存布局和大小上的根本差异。通过引入双描述符视图、物理与逻辑块桥接以及3-描述符卷积传输等机制，系统实现了零开销、无暂存缓冲区且无需数据洗牌的高效RDMA传输。实验证明，该架构在高并发场景下显著提升了吞吐量表现。

未来工作展望：

  1. Mamba1模型支持：目前的3-描述符卷积传输仅支持Mamba2。由于Mamba1的SSM时间形状为$(intermediate\_size // tp, state\_size)$，不允许重建卷积分解所需的$intermediate\_size$。此外，GDN支持（Qwen3.5+）已列入分离式服务的路线图中。
  2. 投机解码（Speculative decoding）：SSM状态传输与投机解码之间的交互尚未得到广泛验证。
  3. 启用HMA时的混合块大小：目前在启用HMA时，尚不支持Prefill和Decode之间存在不同的块大小（即$block\_size\_ratio > 1$）。

补充细节

启动配置指南：
要运行具有分离式P/D的混合SSM模型，需要设置特定的环境变量和启动参数。特别是，异构TP必须设置DS卷积状态布局VLLM_SSM_CONV_STATE_LAYOUT=DS（非异构情况下非必须）。Prefill实例的启动代码示例如下：

# Prefill instance
VLLM_SSM_CONV_STATE_LAYOUT=DS vllm serve nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.85 \
    --trust-remote-code \
    --max-model-len 8192 \
    --block-size 128 \
    --no-disable-hybrid-kv-cache-manager \
    --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'