vLLM v0.20.0 adds disaggregated serving for hybrid SSM models

Introduction

vLLM extended its NIXL-based KV connector to support disaggregated prefill/decode for hybrid SSM-FA models without changing the existing workflow for standard transformers. Hybrid architectures such as NVIDIA Nemotron-H interleave Mamba-style SSM layers with full-attention layers, combining linear-time SSM efficiency with attention expressiveness. The existing NIXL disaggregated P/D design assumed a uniform KV cache format, which does not hold for hybrid models because FA and SSM layers store state in different layouts and sizes.

Background: NIXL KV Transfer Workflow

For standard transformer models, NIXL disaggregated P/D works by registering memory regions, creating per-block descriptors, performing a handshake, and transferring blocks via RDMA. Each worker registers its KV cache tensors with NIXL, then creates descriptors that specify (address, length, device_id) for each block. A handshake exchanges metadata once per prefetch-decode pair. The scheduler tells the decode worker which block IDs to pull; the decode worker maps block ID to descriptor ID, issues an RDMA READ, and polls for completion. With M registered regions and N blocks per region, descriptor index is r * N + b.

The Challenge: FA and SSM State Are Fundamentally Different

Hybrid models break the uniform descriptor assumption because FA layers store per-token K/V pairs while SSM layers store a fixed-size conv state and SSM state. FA layers have KV cache shape [num_blocks, 2, block_size, num_kv_heads, head_dim] (or a variant). SSM layers maintain conv state (conv_dim, state_len) and SSM state (num_heads, head_dim, state_size). These states are not token-based, so the notion of block size differs: each SSM block is a full state snapshot. HMA pools memory across FA and SSM groups, padding SSM rows so both views share the same physical tensor and page size in bytes. Consequently, a single descriptor list with uniform (address, length) cannot correctly index both FA and SSM views.

Dual Descriptor Views

The solution registers two separate descriptor lists over the same physical memory: one list for FA descriptors and another for SSM descriptors, concatenated under a single NIXL transfer handle. FA descriptors occupy the first num_descs = M * N_phys slots, indexing K and V separately for each region. SSM descriptors follow, with each Mamba layer represented by four sub-descriptors (x, B, C, SSM) per block. Block ID to descriptor ID mapping uses region * N_phys + block_id for FA groups and mamba_region_id * N_log + block_id + num_descs for SSM groups.

Physical vs. Logical Block Sizes

FA layers may require a different physical block size for attention kernels, while SSM layers use logical blocks directly, leading to different block counts in the FA and SSM descriptor sections. The ratio logical_block_size / kernel_block_size computes physical_blocks = logical_blocks * ratio for FA layers only. SSM layers always use logical_blocks. This is tracked via _physical_blocks_per_logical, which may differ between prefetch and decode instances when their tensor-parallel sizes differ. The block-ID-to-descriptor-ID mapping uses the appropriate stride depending on whether the group is FA or SSM.

The 3-Descriptors Conv Transfer

For heterogeneous tensor-parallel setups, the conv state is laid out in DS format so each decode rank can read its contiguous x, B, and C slices via three descriptor regions, enabling zero-copy RDMA without extra buffers or reshuffling. The DS layout (dim, state_len) makes each sub-projection contiguous in memory: [x][B][C][SSM]. A decode worker can issue three separate contiguous reads within a single NIXL READ operation to obtain its shard of the conv state. This avoids allocating a staging buffer on the decode side, eliminates post-transfer reshuffling, transfers only the owned 1/TP share of the conv state, and skips HMA padding bytes. The approach works for both homogeneous and heterogeneous TP, with the latter using four descriptor regions per Mamba layer (x, B, C, SSM).

Putting It Together: Nemotron-H Example

A concrete example shows how descriptor registration and transfer work for nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 served with disaggregated P/D at TP=2. The model has 52 layers alternating Mamba and FA, grouped by HMA into 5 groups (4 Mamba, 1 FA) yielding 6 shared KV cache tensors. FA layers use [num_blocks, 2, block_size=400, 4, 128]; SSM layers use [num_blocks, 3, 3072] (conv) plus [num_blocks, 48, 64, 128] (ssm). After HMA padding, both views share the same page size in bytes. The prefetch instance registers the 6 tensors, creates FA descriptors for all regions x N_phys blocks (K and V separate), and appends Mamba descriptors with four sub-regions per block. The decode instance maps block IDs to descriptor IDs using the appropriate stride, issues a single prepopulated READ with both FA and Mamba descriptors, and polls for completion. No intermediate buffers or data reshuffling are needed.

Performance

Disaggregated P/D for hybrid SSM models matches or exceeds co-located serving throughput at high concurrency by isolating decode from prefill interference. Benchmarks on 8x H200 GPUs using NVLink compared a co-located baseline (single instance, TP=8) with a disaggregated setup (1 prefetch instance TP=4 + 1 decode instance TP=4, same total GPU count) on nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8. Sweeping concurrency from 8 to 256 users, the disaggregated Pareto curve dominates the co-located curve at higher batch sizes, showing higher output tokens per second per GPU when decode is separated from prefetch.

Getting Started

To run a hybrid SSM model with disaggregated P/D, set VLLM_SSM_CONV_STATE_LAYOUT=DS and launch prefetch and decode instances with the appropriate arguments. Example command for the prefetch 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"}' The decode instance uses a similar command with kv_role set appropriately (not shown in source). The DS layout is required for heterogeneous TP but optional otherwise.

Limitations and Future Work

Current support covers Mamba2 only; Mamba1 models and GDN layers are not yet supported, speculative decoding interaction has not been extensively validated, and mixed block sizes with HMA are not yet supported when HMA is enabled. Mamba1’s SSM temporal shape prevents reconstruction of intermediate_size needed for conv decomposition. GDN support is listed in the disaggregated roadmap. Speculative decoding and block-size ratio handling with HMA remain open work.

Acknowledgments

Thanks to Thomas Parnell (IBM Research) and Roi Koren (NVIDIA) for contributions.