vLLM v0.20.0 为混合 SSM 模型添加了分离式服务

介绍

vLLM 扩展了基于 NIXL 的 KV 连接器，以支持混合 SSM-FA 模型的分离式预填充/解码，而无需更改标准 Transformer 的现有工作流。

诸如 NVIDIA Nemotron-H 等混合架构交错使用 Mamba 风格的 SSM 层和全注意力层，将线性时间的 SSM 效率与注意力的表达能力相结合。

现有的 NIXL 分离式 P/D 设计假设 KV 缓存格式统一，但对混合模型并不适用，因为 FA 层和 SSM 层以不同的布局和大小存储状态。

背景：NIXL KV 传输工作流

对于标准 Transformer 模型，NIXL 分离式 P/D 通过注册内存区域、创建每块描述符、进行握手并通过 RDMA 传输块来工作。

每个工作节点使用 NIXL 注册其 KV 缓存张量，然后创建描述符，指定每块的 (address, length, device_id)。握手在每个预取‑解码对之间交换一次元数据。调度器告知解码工作节点需要拉取的块 ID；解码工作节点将块 ID 映射到描述符 ID，发起 RDMA READ，并轮询完成情况。若有 M 个已注册区域且每个区域有 N 块，则描述符索引为 r * N + b。

挑战：FA 与 SSM 状态本质不同

混合模型打破了统一描述符的假设，因为 FA 层存储每个 token 的 K/V 对，而 SSM 层存储固定大小的卷积状态和 SSM 状态。

FA 层的 KV 缓存形状为 [num_blocks, 2, block_size, num_kv_heads, head_dim]（或其变体）。SSM 层维护卷积状态 (conv_dim, state_len) 和 SSM 状态 (num_heads, head_dim, state_size)。这些状态并非基于 token，因此块大小的概念不同：每个 SSM 块是完整的状态快照。HMA 在 FA 与 SSM 组之间统一内存池，通过填充 SSM 行使两种视图共享相同的物理张量和字节页大小。因此，使用统一 (address, length) 的单一描述符列表无法正确索引 FA 与 SSM 两种视图。

双描述符视图

解决方案是在同一块物理内存上注册两个独立的描述符列表：一个用于 FA 描述符，另一个用于 SSM 描述符，并在单个 NIXL 传输句柄下进行拼接。

FA 描述符占据前 num_descs = M * N_phys 个槽位，对每个区域的 K 与 V 分别建立索引。随后是 SSM 描述符，每个 Mamba 层在每块上由四个子描述符 (x, B, C, SSM) 表示。块 ID 到描述符 ID 的映射对 FA 组使用 region * N_phys + block_id，而对 SSM 组使用 mamba_region_id * N_log + block_id + num_descs。

物理块大小 vs. 逻辑块大小

FA 层可能需要不同的物理块大小以适配注意力内核，而 SSM 层直接使用逻辑块，这导致 FA 与 SSM 描述符部分的块数量不同。

比例 logical_block_size / kernel_block_size 用于计算 FA 层的 physical_blocks = logical_blocks * ratio。SSM 层始终使用 logical_blocks。此信息通过 _physical_blocks_per_logical 记录，当预取与解码实例的张量并行大小不同时时可能会有差异。块 ID 到描述符 ID 的映射会根据组是 FA 还是 SSM 使用相应的步幅。

三描述符卷积传输

对于异构张量并行配置，卷积状态采用 DS 格式布局，使每个解码 rank 能够通过三个描述符区域读取其连续的 x、B、C 切片，从而实现零拷贝 RDMA，无需额外缓冲区或重新排列。

DS 布局 (dim, state_len) 使每个子投影在内存中连续排列：[x][B][C][SSM]。解码工作节点可以在一次 NIXL READ 操作中发起三次连续读取，以获取其对应的卷积状态分片。这避免了在解码端分配暂存缓冲区，消除了传输后重新排列，仅传输拥有的 1/TP 的卷积状态份额，并跳过 HMA 填充字节。该方法适用于同构和异构 TP，后者在每个 Mamba 层使用四个描述符区域 (x, B, C, SSM)。

综合示例：Nemotron-H

下面的具体示例展示了在 TP=2 时，使用分离式 P/D 为 nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 提供服务时，描述符注册与传输的工作方式。

该模型有 52 层，交替出现 Mamba 和 FA 层，经过 HMA 分组为 5 组（4 个 Mamba 组，1 个 FA 组），产生 6 个共享的 KV 缓存张量。FA 层使用形状 [num_blocks, 2, block_size=400, 4, 128]；SSM 层使用 [num_blocks, 3, 3072]（卷积）以及 [num_blocks, 48, 64, 128]（ssm）。经过 HMA 填充后，两种视图共享相同的字节页大小。预取实例注册这 6 个张量，为所有区域的 N_phys 块创建 FA 描述符（K 与 V 分开），并为每块追加四个子区域的 Mamba 描述符。解码实例使用相应的步幅将块 ID 映射到描述符 ID，发起一次包含 FA 与 Mamba 描述符的预填充 READ，并轮询完成情况。无需中间缓冲区或数据重新排列。

性能

通过将解码与预填充的干扰隔离，混合 SSM 模型的分离式 P/D 在高并发下的吞吐量可与或超过同机服务的水平。

在 8 台 H200 GPU（使用 NVLink）上进行基准测试，将同机基线（单实例，TP=8）与分离式配置（1 个预取实例 TP=4 + 1 个解码实例 TP=4，总 GPU 数相同）在 nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 上进行比较。并发用户数从 8 到 256 逐步增加，分离式的帕累托曲线在更大批次时超越同机曲线，表明当解码与预取分离时，每 GPU 的每秒输出 token 数更高。

入门指南

要运行支持分离式 P/D 的混合 SSM 模型，请设置 VLLM_SSM_CONV_STATE_LAYOUT=DS 并使用相应参数启动预取和解码实例。

示例预取实例命令：

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"}'

解码实例使用类似的命令，只需将 kv_role 设置为相应值（源码中未展示）。DS 布局在异构 TP 场景下是必需的，其他情况下可选。

限制与未来工作

当前仅支持 Mamba2；Mamba1 模型和 GDN 层尚未支持，推测解码的交互也未得到充分验证，并且在启用 HMA 时尚不支持混合块大小。

Mamba1 的 SSM 时间形状阻碍了用于卷积分解的 intermediate_size 重建。GDN 支持已列入分离式路线图。推测解码以及在 HMA 环境下的块大小比例处理仍待进一步工作。

致谢

感谢 Thomas Parnell（IBM Research）和 Roi Koren（NVIDIA）的贡献。