[NIXL][Mamba] Add Mamba1 support to NIXL P/D disaggregation (#45019)
Signed-off-by: Josephasafg <ajgard7@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Asaf Gardin committed
Jun 25, 2026 at 15:50 UTC
15be78732bac03af62dc05d588e245271b72f10d
5 files changed
+104
-35
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+2
@@ -28,6 +28,8 @@ hybrid_ssm_configs=(
28
# GDN (Qwen3.5)
29
"VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
30
"VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
31
+ # Mamba1 (Jamba)
32
+ "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ai21labs/AI21-Jamba2-3B VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
33
)
34
sw_attn_configs=(
35
# NOTE: gemma3 does not work with FlashInfer
tests/v1/kv_connector/nixl_integration/test_accuracy.py
+1
@@ -25,6 +25,7 @@ EXPECTED_VALUES = {
25
"ibm-granite/granite-4.0-h-tiny": 0.77,
26
"Qwen/Qwen3.5-0.8B": 0.33,
27
"google/gemma-4-E2B-it": 0.485,
28
+ "ai21labs/AI21-Jamba2-3B": 0.74,
29
}
30
31
SIMPLE_PROMPT = (
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+45
@@ -630,6 +630,19 @@ def _make_mock_worker_for_desc_ids(
630
worker._has_mamba = has_mamba
631
worker._group_spec_types = group_spec_types
632
worker.block_len_per_layer = block_len_per_layer or [100]
633
+ worker._conv_decomp = None
634
+ if has_mamba:
635
+ from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( # noqa: E501
636
+ MambaConvSplitInfo,
637
+ )
638
+
639
+ # Mamba2/GDN layout: 3 conv sub-projections -> 4 NIXL regions per layer.
640
+ worker._conv_decomp = MambaConvSplitInfo(
641
+ conv_rows=3,
642
+ local_proj_dims=(1, 1, 1),
643
+ conv_dtype_size=2,
644
+ ssm_sizes=(0, 0),
645
+ )
646
worker._compute_desc_ids = NixlConnectorWorker._compute_desc_ids.__get__(
647
worker, NixlConnectorWorker
648
)
@@ -976,6 +989,37 @@ def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_rati
989
(256, 256, 768),
990
id="qwen35_27b_tp8",
991
),
992
+ # ai21labs/AI21-Jamba2-Mini (Mamba1)
993
+ # mamba d_inner = mamba_expand(2) * hidden_size(4096) = 8192
994
+ # mamba_d_state=16, mamba_d_conv=4 → conv_rows=3.
995
+ # Conv state holds only x: a single contiguous sub-projection.
996
+ pytest.param(
997
+ "mamba1",
998
+ 1,
999
+ 8192,
1000
+ 3,
1001
+ (8192, 16),
1002
+ (8192,),
1003
+ id="jamba_mini_tp1",
1004
+ ),
1005
+ pytest.param(
1006
+ "mamba1",
1007
+ 4,
1008
+ 2048,
1009
+ 3,
1010
+ (2048, 16),
1011
+ (2048,),
1012
+ id="jamba_mini_tp4",
1013
+ ),
1014
+ pytest.param(
1015
+ "mamba1",
1016
+ 8,
1017
+ 1024,
1018
+ 3,
1019
+ (1024, 16),
1020
+ (1024,),
1021
+ id="jamba_mini_tp8",
1022
+ ),
1023
],
1024
)
1025
def test_derive_mamba_conv_split(
@@ -999,6 +1043,7 @@ def test_derive_mamba_conv_split(
1043
from vllm.v1.kv_cache_interface import MambaSpec
1044
1045
_TYPE_MAP = {
1046
+ "mamba1": MambaAttentionBackendEnum.MAMBA1,
1047
"mamba2": MambaAttentionBackendEnum.MAMBA2,
1048
"gdn_attention": MambaAttentionBackendEnum.GDN_ATTN,
1049
}
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+14
-8
@@ -96,7 +96,13 @@ class NixlBaseConnectorWorker:
96
) -> np.ndarray:
97
"""Compute NIXL descriptor IDs for given block IDs."""
98
num_fa_regions = self.num_regions
99
- num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0
99
+ num_ssm_regions = 0
100
+ if self._has_mamba:
101
+ assert self._conv_decomp is not None
102
+ # NIXL regions per SSM layer = conv sub-projections + 1 SSM temporal
103
+ # (Mamba2/GDN: 3+1=4; Mamba1: 1+1=2).
104
+ ssm_regions_per_layer = len(self._conv_decomp.local_conv_offsets) + 1
105
+ num_ssm_regions = len(self.block_len_per_layer) * ssm_regions_per_layer
106
107
num_blocks = dst_num_blocks
108
if block_size_ratio is not None:
@@ -279,8 +285,8 @@ class NixlBaseConnectorWorker:
285
# ---- Model state (derived from model config) ----
286
mamba_ssm_size = (0, 0)
287
# Conv state sub-projection decomposition (None when no Mamba).
282
- # The 3-read transfer requires DS (dim, state_len) conv layout so
283
- # that x/B/C sub-projections are contiguous in memory.
288
+ # The transfer requires DS (dim, state_len) conv layout so that
289
+ # conv sub-projections are contiguous in memory.
290
self._conv_decomp: MambaConvSplitInfo | None = None
291
self._has_mamba = any(
292
isinstance(g.kv_cache_spec, MambaSpec)
@@ -1186,8 +1192,8 @@ class NixlBaseConnectorWorker:
1192
base_addresses: list[int],
1193
block_size_ratio: int,
1194
) -> list[tuple[int, int, int]]:
1189
- """Build 4 desc regions (x, B, C, ssm) per layer for local mamba
1190
- blocks, enabling the 3-read transfer with DS conv layout."""
1195
+ """Build desc regions (conv sub-projections + ssm) per layer for
1196
+ local mamba blocks with DS conv layout."""
1197
assert block_size_ratio == 1, (
1198
"Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
1199
f"Got block_size_ratio={block_size_ratio}."
@@ -1227,9 +1233,9 @@ class NixlBaseConnectorWorker:
1233
tp_ratio: int,
1234
transfer_info: EngineTransferInfo,
1235
) -> list[tuple[int, int, int]]:
1230
- """Build 4 remote desc regions (proj0, proj1, proj2, ssm) per layer
1231
- for the 3-read transfer. For hetero-TP, each D rank reads only its
1232
- sub-projection slice from the P rank."""
1236
+ """Build remote desc regions (conv sub-projections + ssm) per layer.
1237
+ For hetero-TP, each D rank reads only its sub-projection slice from
1238
+ the P rank."""
1239
assert self._conv_decomp is not None
1240
effective_ratio = max(tp_ratio, 1)
1241
# Mamba conv state is always TP-sharded, even when attention KV
vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
+42
-27
@@ -1,12 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Mamba conv-state sub-projection decomposition for the 3-read transfer.
3
+"""Mamba conv-state sub-projection decomposition for NIXL transfer.
4
5
With DS conv state layout (dim, state_len), sub-projections are
6
-contiguous in memory. Each D rank reads its slices via 3 separate
6
+contiguous in memory. Each D rank reads its slices via separate
7
RDMA transfers — no P-side permutation needed.
8
9
Supported model types:
10
+ - Mamba1: conv = [x], temporal = (intermediate_size, state_size)
11
- Mamba2: conv = [x, B, C], temporal = (num_heads, head_dim)
12
- GDN (Gated Delta Net): conv = [Q, K, V] (dim(Q)==dim(K)),
13
temporal = (num_v_heads, v_dim, k_dim)
@@ -24,18 +25,21 @@ from vllm.v1.kv_cache_interface import MambaSpec
25
26
@dataclass(frozen=True)
27
class MambaConvSplitInfo:
27
- """Per-rank byte sizes of the 3 conv sub-projections.
28
+ """Per-rank byte sizes of the conv sub-projections.
29
30
Used by both P and D sides for NIXL descriptor registration.
31
All fields are LOCAL to this engine's TP (already divided by TP size).
32
33
DS memory layout within one page (contiguous):
34
+ Mamba1: |---- x ----| (single sub-projection, no decomposition)
35
Mamba2: |-- x --|- B -|- C -| (B == C)
36
GDN: |- Q -|- K -|-- V --| (dim(Q)==dim(K), V may differ)
37
"""
38
39
conv_rows: int # conv_kernel - 1 (typically 3)
38
- local_proj_dims: tuple[int, int, int] # per-rank column counts per sub-proj
40
+ # Per-rank column counts per sub-projection:
41
+ # 1 entry for Mamba1, 3 for Mamba2/GDN.
42
+ local_proj_dims: tuple[int, ...]
43
conv_dtype_size: int # bytes per element (e.g. 2 for float16)
44
ssm_sizes: tuple[int, int] # (conv_state_bytes, ssm_state_bytes)
45
@@ -45,10 +49,10 @@ class MambaConvSplitInfo:
49
return sum(self.local_proj_dims)
50
51
@property
48
- def proj_bytes(self) -> tuple[int, int, int]:
49
- """Byte sizes of the 3 sub-projections for one rank."""
52
+ def proj_bytes(self) -> tuple[int, ...]:
53
+ """Byte sizes of the sub-projections for one rank."""
54
row_bytes = self.conv_rows * self.conv_dtype_size
51
- return tuple(d * row_bytes for d in self.local_proj_dims) # type: ignore[return-value]
55
+ return tuple(d * row_bytes for d in self.local_proj_dims)
56
57
@property
58
def local_conv_offsets(self) -> list[tuple[int, int]]:
@@ -57,8 +61,12 @@ class MambaConvSplitInfo:
61
62
Used by both P and D for local descriptor registration.
63
"""
60
- conv0, conv1, conv2 = self.proj_bytes
61
- return [(0, conv0), (conv0, conv1), (conv0 + conv1, conv2)]
64
+ offsets: list[tuple[int, int]] = []
65
+ offset = 0
66
+ for size in self.proj_bytes:
67
+ offsets.append((offset, size))
68
+ offset += size
69
+ return offsets
70
71
def remote_conv_offsets(
72
self, local_rank_offset: int, tp_ratio: int
@@ -76,28 +84,23 @@ class MambaConvSplitInfo:
84
P page. Local dims are scaled down by |tp_ratio|
85
to get P-sized offsets.
86
"""
79
- conv0, conv1, conv2 = self.proj_bytes
87
+ offsets: list[tuple[int, int]] = []
88
if tp_ratio >= 1:
81
- remote_conv0 = conv0 * tp_ratio
82
- remote_conv1 = conv1 * tp_ratio
83
- return [
84
- (local_rank_offset * conv0, conv0),
85
- (remote_conv0 + local_rank_offset * conv1, conv1),
86
- (remote_conv0 + remote_conv1 + local_rank_offset * conv2, conv2),
87
- ]
89
+ remote_base = 0
90
+ for size in self.proj_bytes:
91
+ offsets.append((remote_base + local_rank_offset * size, size))
92
+ remote_base += size * tp_ratio
93
else:
94
# NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages
95
# are smaller than D's. Local dims are D-sized, but we need
96
# P-sized offsets. Scale down by |tp_ratio|.
97
abs_ratio = -tp_ratio
93
- remote_conv0 = conv0 // abs_ratio
94
- remote_conv1 = conv1 // abs_ratio
95
- remote_conv2 = conv2 // abs_ratio
96
- return [
97
- (0, remote_conv0),
98
- (remote_conv0, remote_conv1),
99
- (remote_conv0 + remote_conv1, remote_conv2),
100
- ]
98
+ remote_base = 0
99
+ for size in self.proj_bytes:
100
+ remote_size = size // abs_ratio
101
+ offsets.append((remote_base, remote_size))
102
+ remote_base += remote_size
103
+ return offsets
104
105
106
def derive_mamba_conv_split(
@@ -120,12 +123,13 @@ def derive_mamba_conv_split(
123
conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes).
124
"""
125
_supported = (
126
+ MambaAttentionBackendEnum.MAMBA1,
127
MambaAttentionBackendEnum.MAMBA2,
128
MambaAttentionBackendEnum.GDN_ATTN,
129
)
130
if mamba_spec.mamba_type not in _supported:
131
raise NotImplementedError(
128
- f"3-read conv transfer only supports Mamba2 and GDN models, "
132
+ f"Conv transfer only supports Mamba1, Mamba2 and GDN models, "
133
f"got mamba_type={mamba_spec.mamba_type!r}."
134
)
135
@@ -149,7 +153,18 @@ def derive_mamba_conv_split(
153
conv_state_bytes = torch.Size(mamba_spec.shapes[0]).numel() * conv_dtype_size
154
ssm_state_bytes = torch.Size(mamba_spec.shapes[1]).numel() * ssm_dtype_size
155
152
- if mamba_spec.mamba_type == MambaAttentionBackendEnum.MAMBA2:
156
+ local_proj_dims: tuple[int, ...]
157
+ if mamba_spec.mamba_type == MambaAttentionBackendEnum.MAMBA1:
158
+ # Mamba1 conv state holds only x (no B/C), so it's a single
159
+ # contiguous TP shard with no sub-projection decomposition.
160
+ temporal_shape = mamba_spec.shapes[1]
161
+ assert temporal_shape[0] == local_conv_dim, (
162
+ f"Mamba1 temporal state dim ({temporal_shape[0]}) doesn't match "
163
+ f"conv dim ({local_conv_dim}); both should be "
164
+ f"intermediate_size/TP."
165
+ )
166
+ local_proj_dims = (local_conv_dim,)
167
+ elif mamba_spec.mamba_type == MambaAttentionBackendEnum.MAMBA2:
168
# NOTE (ZhanqiuHu): intermediate_size (= global x dim) is not stored
169
# in MambaSpec, so we reconstruct it from the SSM temporal state shape:
170
# shapes[1] = (local_num_heads, head_dim), already divided by TP.