vLLM v0.20.0 為混合 SSM 模型新增分散式服務

vLLM v0.20.0 為混合 SSM 模型新增分散式服務

vLLM v0.20.0 為混合 SSM-FA 模型引入分散式 prefill/decode 服務，透過雙描述子視圖與三描述子卷積狀態傳輸，實現高效的 KV 傳輸。

vLLM v0.20.0 為混合 SSM 模型新增分散式服務

介紹

vLLM 擴充了其基於 NIXL 的 KV 連接器，以支援混合 SSM-FA 模型的分散式 prefill/decode，且不需更改標準 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)。握手在每對 prefetch‑decode 之間交換一次中繼資料。排程器告訴 decode 工作者要拉取哪些 block ID；decode 工作者將 block ID 映射到 descriptor 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) 代表。Block ID 到 descriptor ID 的映射對於 FA 群組使用 region * N_phys + block_id，對於 SSM 群組則使用 mamba_region_id * N_log + block_id + num_descs。

實體與邏輯區塊大小

FA 層可能需要不同的實體區塊大小以配合注意力核心，而 SSM 層直接使用邏輯區塊，導致 FA 與 SSM 描述子區段的區塊數量不同。

比例 logical_block_size / kernel_block_size 用於計算 FA 層的 physical_blocks = logical_blocks * ratio。SSM 層始終使用 logical_blocks。此資訊透過 _physical_blocks_per_logical 追蹤，當 prefetch 與 decode 實例的 tensor‑parallel 大小不同時可能會有差異。block‑ID 到 descriptor‑ID 的映射會根據群組是 FA 或 SSM 而使用相應的步幅。

三描述子卷積傳輸

對於異質 tensor‑parallel 配置，卷積狀態以 DS 格式排列，使每個 decode rank 能透過三個描述子區域讀取其連續的 x、B、C 切片，從而在不使用額外緩衝或重新排列的情況下實現零拷貝 RDMA。

DS 佈局 (dim, state_len) 使每個子投影在記憶體中連續排列：[x][B][C][SSM]。decode 工作者可在單一 NIXL READ 操作內發出三個分離的連續讀取，以取得其卷積狀態的分片。此方式避免在 decode 端分配暫存緩衝、消除傳輸後的重新排列、僅傳輸其擁有的 1/TP 卷積狀態份額，並跳過 HMA 的填充位元組。此方法同時適用於同質與異質 TP，後者在每個 Mamba 層使用四個描述子區域 (x, B, C, SSM)。

整合示例：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 快取張量。FA 層使用 [num_blocks, 2, block_size=400, 4, 128]；SSM 層使用 [num_blocks, 3, 3072]（卷積）加上 [num_blocks, 48, 64, 128]（ssm）。經過 HMA 填充後，兩種視圖共享相同的頁面大小（位元組）。prefetch 實例註冊這 6 個張量，為所有 region × N_phys 區塊建立 FA 描述子（K 與 V 分開），並為每個區塊附加四個子區域的 Mamba 描述子。decode 實例使用相應的步幅將 block ID 映射到 descriptor ID，發出一次包含 FA 與 Mamba 描述子的預先填充 READ，並輪詢完成。無需中介緩衝或資料重新排列。

效能

對於混合 SSM 模型的分散式 P/D，透過將 decode 與 prefill 干擾分離，在高併發情況下的吞吐量可匹配或超過同地服務的表現。

在使用 NVLink 的 8 × H200 GPU 上進行基準測試，將同地基線（單一實例，TP=8）與分散式配置（1 個 prefetch 實例 TP=4 + 1 個 decode 實例 TP=4，GPU 數量相同）在 nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 上進行比較。併發數從 8 到 256 使用者逐步提升，分散式的 Pareto 曲線在較大批次大小時超越同地曲線，顯示當 decode 與 prefetch 分離時，每 GPU 的每秒輸出 token 數更高。

快速入門

要以分散式 P/D 執行混合 SSM 模型，請設定 VLLM_SSM_CONV_STATE_LAYOUT=DS，並以適當參數啟動 prefetch 與 decode 實例。

prefetch 實例的範例指令：

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

decode 實例使用類似的指令，將 kv_role 設為相應值（此處未示範）。DS 佈局在異質 TP 時是必要的，其他情況則為可選。

限制與未來工作

目前僅支援 Mamba2；Mamba1 模型與 GDN 層尚未支援，推測式解碼的互動尚未廣泛驗證，且在啟用 HMA 時仍不支援混合區塊大小。

Mamba1 的 SSM 時間形狀阻礙了 conv 分解所需的 intermediate_size 重建。GDN 支援已列入分散式路線圖。推測式解碼以及在 HMA 下的區塊大小比例處理仍待完成。

致謝

感謝 Thomas Parnell（IBM Research）與 Roi Koren（NVIDIA）的貢獻。

Sources