@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 24337fb860
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 16:30 UTC

PD disagg with NIXL Connector: GDN support (Qwen3.5) (#41869)

Signed-off-by: Zhanqiu Hu <zhu@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

zhanqiuhu committed May 14, 2026 at 10:33 UTC 24337fb860a89f2087df27b02b7c9620a7fd0093
6 files changed +244 -83
.buildkite/test_areas/disaggregated.yaml
+1 -1
@@ -51,7 +51,7 @@ steps:
51
52 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
53 key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
54 - timeout_in_minutes: 20
54 + timeout_in_minutes: 25
55 working_dir: "/vllm-workspace/tests"
56 num_devices: 4
57 source_file_dependencies:
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+3
@@ -21,6 +21,9 @@ hybrid_ssm_configs=(
21 "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
22 # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
23 "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
24 + # GDN (Qwen3.5)
25 + "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
26 + "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--no-async-scheduling"
27 )
28 sw_attn_configs=(
29 # NOTE: gemma3 does not work with FlashInfer
tests/v1/kv_connector/nixl_integration/test_accuracy.py
+1
@@ -20,6 +20,7 @@ EXPECTED_VALUES = {
20 "google/gemma-3-4b-it": 0.74,
21 "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84,
22 "ibm-granite/granite-4.0-h-tiny": 0.80,
23 + "Qwen/Qwen3.5-0.8B": 0.33,
24 }
25
26 SIMPLE_PROMPT = (
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+126
@@ -742,6 +742,132 @@ def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_rati
742 assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio
743
744
745 +@pytest.mark.cpu_test
746 +@pytest.mark.parametrize(
747 + "mamba_type,local_tp,conv_dim_local,conv_rows,temporal_shape,expected_proj_dims",
748 + [
749 + # nvidia/Nemotron-H-8B-Base-8K (Mamba2)
750 + # mamba_num_heads=128, head_dim=64, n_groups=8, ssm_state_size=128
751 + pytest.param(
752 + "mamba2",
753 + 1,
754 + 10240,
755 + 3,
756 + (128, 64, 128),
757 + (8192, 1024, 1024),
758 + id="nemotron_h_8b_tp1",
759 + ),
760 + pytest.param(
761 + "mamba2",
762 + 4,
763 + 2560,
764 + 3,
765 + (32, 64, 128),
766 + (2048, 256, 256),
767 + id="nemotron_h_8b_tp4",
768 + ),
769 + # nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B (Mamba2)
770 + # mamba_num_heads=64, head_dim=64, n_groups=8, ssm_state_size=128
771 + pytest.param(
772 + "mamba2",
773 + 1,
774 + 6144,
775 + 3,
776 + (64, 64, 128),
777 + (4096, 1024, 1024),
778 + id="nemotron_nano_30b_tp1",
779 + ),
780 + # Qwen/Qwen3.5-0.8B (GDN, symmetric: num_v=num_k=16)
781 + # key_dim=2048, value_dim=2048, conv_dim=6144
782 + pytest.param(
783 + "gdn_attention",
784 + 1,
785 + 6144,
786 + 3,
787 + (16, 128, 128),
788 + (2048, 2048, 2048),
789 + id="qwen35_08b_tp1",
790 + ),
791 + pytest.param(
792 + "gdn_attention",
793 + 4,
794 + 1536,
795 + 3,
796 + (4, 128, 128),
797 + (512, 512, 512),
798 + id="qwen35_08b_tp4",
799 + ),
800 + # Qwen/Qwen3.5-4B (GDN, asymmetric: num_v=32, num_k=16, K:V=1:2)
801 + # key_dim=2048, value_dim=4096, conv_dim=8192
802 + pytest.param(
803 + "gdn_attention",
804 + 1,
805 + 8192,
806 + 3,
807 + (32, 128, 128),
808 + (2048, 2048, 4096),
809 + id="qwen35_4b_tp1",
810 + ),
811 + # Qwen/Qwen3.5-27B (GDN, asymmetric: num_v=48, num_k=16, K:V=1:3)
812 + # key_dim=2048, value_dim=6144, conv_dim=10240
813 + pytest.param(
814 + "gdn_attention",
815 + 1,
816 + 10240,
817 + 3,
818 + (48, 128, 128),
819 + (2048, 2048, 6144),
820 + id="qwen35_27b_tp1",
821 + ),
822 + pytest.param(
823 + "gdn_attention",
824 + 8,
825 + 1280,
826 + 3,
827 + (6, 128, 128),
828 + (256, 256, 768),
829 + id="qwen35_27b_tp8",
830 + ),
831 + ],
832 +)
833 +def test_derive_mamba_conv_split(
834 + monkeypatch,
835 + mamba_type,
836 + local_tp,
837 + conv_dim_local,
838 + conv_rows,
839 + temporal_shape,
840 + expected_proj_dims,
841 +):
842 + """Parametrized test for derive_mamba_conv_split with real model configs.
843 +
844 + Values generated by verify_conv_split.py which loads HuggingFace configs
845 + and calls vLLM's derive_mamba_conv_split directly.
846 + """
847 + from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
848 + derive_mamba_conv_split,
849 + )
850 + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
851 + from vllm.v1.kv_cache_interface import MambaSpec
852 +
853 + _TYPE_MAP = {
854 + "mamba2": MambaAttentionBackendEnum.MAMBA2,
855 + "gdn_attention": MambaAttentionBackendEnum.GDN_ATTN,
856 + }
857 + mamba_type_enum = _TYPE_MAP[mamba_type]
858 +
859 + monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "DS")
860 + spec = MambaSpec(
861 + block_size=64,
862 + shapes=((conv_dim_local, conv_rows), temporal_shape),
863 + dtypes=(torch.bfloat16, torch.bfloat16),
864 + mamba_type=mamba_type_enum,
865 + )
866 + out = derive_mamba_conv_split(spec, local_tp=local_tp)
867 + assert out.local_proj_dims == expected_proj_dims
868 + assert out.conv_rows == conv_rows
869 +
870 +
871 @pytest.mark.cpu_test
872 @pytest.mark.parametrize(
873 "mamba_enabled,swa_enabled,"
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+3 -13
@@ -1057,8 +1057,8 @@ class NixlConnectorWorker:
1057 tp_ratio: int,
1058 transfer_info: EngineTransferInfo,
1059 ) -> list[tuple[int, int, int]]:
1060 - """Build 4 remote desc regions (x, B, C, ssm) per layer for
1061 - the 3-read transfer. For hetero-TP, each D rank reads only its
1060 + """Build 4 remote desc regions (proj0, proj1, proj2, ssm) per layer
1061 + for the 3-read transfer. For hetero-TP, each D rank reads only its
1062 sub-projection slice from the P rank."""
1063 assert self._conv_decomp is not None
1064 effective_ratio = max(tp_ratio, 1)
@@ -1067,20 +1067,10 @@ class NixlConnectorWorker:
1067 local_offset = self.tp_rank % effective_ratio
1068 conv_size_remote = nixl_agent_meta.ssm_sizes[0]
1069
1070 + conv_offsets = self._conv_decomp.remote_conv_offsets(local_offset, tp_ratio)
1071 if tp_ratio >= 1:
1071 - # D_TP >= P_TP: P page is larger, D reads its slice.
1072 - conv_offsets = self._conv_decomp.remote_conv_offsets(
1073 - local_offset, effective_ratio
1074 - )
1072 ssm_read_size = self._mamba_ssm_size[1]
1073 else:
1077 - # NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages
1078 - # are smaller than D's. self._conv_decomp has D-sized dimensions,
1079 - # but we need P-sized offsets. Scale down by |tp_ratio|.
1080 - abs_ratio = -tp_ratio
1081 - xb_p = self._conv_decomp.x_bytes // abs_ratio
1082 - bb_p = self._conv_decomp.b_bytes // abs_ratio
1083 - conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)]
1074 ssm_read_size = nixl_agent_meta.ssm_sizes[1]
1075
1076 remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical
vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
+110 -69
@@ -2,9 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Mamba conv-state sub-projection decomposition for the 3-read transfer.
4
5 -With DS conv state layout (dim, state_len), x/B/C sub-projections are
6 -contiguous in memory. Each D rank reads its x, B, C slices via 3
7 -separate RDMA transfers — no P-side permutation needed.
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
7 +RDMA transfers — no P-side permutation needed.
8 +
9 +Supported model types:
10 + - Mamba2: conv = [x, B, C], temporal = (num_heads, head_dim)
11 + - GDN (Gated Delta Net): conv = [Q, K, V] (dim(Q)==dim(K)),
12 + temporal = (num_v_heads, v_dim, k_dim)
13 """
14
15 import math
@@ -19,97 +24,109 @@ from vllm.v1.kv_cache_interface import MambaSpec
24
25 @dataclass(frozen=True)
26 class MambaConvSplitInfo:
22 - """Per-rank byte sizes of x, B, C sub-projections in the Mamba conv state.
27 + """Per-rank byte sizes of the 3 conv sub-projections.
28
29 Used by both P and D sides for NIXL descriptor registration.
30 All fields are LOCAL to this engine's TP (already divided by TP size).
31
27 - DS memory layout within one page (contiguous in memory):
28 - |--- x (x_local * conv_rows) ---|- B (b_local * conv_rows) -|- C -|
32 + DS memory layout within one page (contiguous):
33 + Mamba2: |-- x --|- B -|- C -| (B == C)
34 + GDN: |- Q -|- K -|-- V --| (dim(Q)==dim(K), V may differ)
35 """
36
37 conv_rows: int # conv_kernel - 1 (typically 3)
32 - x_local: int # intermediate_size / TP (columns for x)
33 - b_local: int # groups_ss / TP (columns for B; C is same size)
38 + local_proj_dims: tuple[int, int, int] # per-rank column counts per sub-proj
39 conv_dtype_size: int # bytes per element (e.g. 2 for float16)
40 ssm_sizes: tuple[int, int] # (conv_state_bytes, ssm_state_bytes)
41
42 @property
38 - def conv_dim_local(self) -> int:
39 - """Total conv columns per rank: x + B + C."""
40 - return self.x_local + 2 * self.b_local
43 + def local_conv_dim(self) -> int:
44 + """Total conv columns per rank."""
45 + return sum(self.local_proj_dims)
46
47 @property
43 - def x_bytes(self) -> int:
44 - """Byte size of the x sub-projection for one rank."""
45 - return self.x_local * self.conv_rows * self.conv_dtype_size
46 -
47 - @property
48 - def b_bytes(self) -> int:
49 - """Byte size of the B (or C) sub-projection for one rank."""
50 - return self.b_local * self.conv_rows * self.conv_dtype_size
48 + def proj_bytes(self) -> tuple[int, int, int]:
49 + """Byte sizes of the 3 sub-projections for one rank."""
50 + 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]
52
53 @property
54 def local_conv_offsets(self) -> list[tuple[int, int]]:
54 - """(byte_offset, byte_size) of x, B, C within this engine's page.
55 + """(byte_offset, byte_size) of each sub-projection within this
56 + engine's page.
57
58 Used by both P and D for local descriptor registration.
59 """
58 - xb = self.x_bytes
59 - bb = self.b_bytes
60 - return [(0, xb), (xb, bb), (xb + bb, bb)]
60 + conv0, conv1, conv2 = self.proj_bytes
61 + return [(0, conv0), (conv0, conv1), (conv0 + conv1, conv2)]
62
63 def remote_conv_offsets(
64 self, local_rank_offset: int, tp_ratio: int
65 ) -> list[tuple[int, int]]:
65 - """(byte_offset, byte_size) of this D rank's x, B, C slice within
66 - one P page.
66 + """(byte_offset, byte_size) of this D rank's sub-projection slices
67 + within one P page.
68
69 Used by D side only, during remote descriptor registration.
70
71 Args:
72 local_rank_offset: which slice this D rank reads.
72 - tp_ratio > 0: tp_rank % tp_ratio (selects slice of P's page).
73 - tp_ratio < 0: always 0 (read P's full page).
74 - tp_ratio: effective ratio (>= 1 when D_TP > P_TP, 1 when
75 - P_TP > D_TP since each P rank is read in full).
73 + tp_ratio: signed TP ratio.
74 + >= 1: D_TP >= P_TP — P page is larger, D reads its slice.
75 + < 0: P_TP > D_TP — P pages are smaller, D reads entire
76 + P page. Local dims are scaled down by |tp_ratio|
77 + to get P-sized offsets.
78 """
77 - xb = self.x_bytes
78 - bb = self.b_bytes
79 - xr = xb * tp_ratio # full remote x section in bytes
80 - br = bb * tp_ratio # full remote B section in bytes
81 - return [
82 - (local_rank_offset * xb, xb),
83 - (xr + local_rank_offset * bb, bb),
84 - (xr + br + local_rank_offset * bb, bb),
85 - ]
79 + conv0, conv1, conv2 = self.proj_bytes
80 + 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 + ]
88 + else:
89 + # NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages
90 + # are smaller than D's. Local dims are D-sized, but we need
91 + # P-sized offsets. Scale down by |tp_ratio|.
92 + 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 + ]
101
102
103 def derive_mamba_conv_split(
104 mamba_spec: MambaSpec,
105 local_tp: int,
106 ) -> MambaConvSplitInfo:
92 - """Derive per-rank x/B/C byte sizes from a MambaSpec.
107 + """Derive per-rank sub-projection byte sizes from a MambaSpec.
108
109 Called once at init on both P and D. Decomposes the conv dimension
95 - (= intermediate_size + 2 * groups_ss) into its x, B, C parts.
110 + into its sub-projection parts based on the model type.
111
112 Args:
113 mamba_spec: MambaSpec whose shapes are:
114 shapes[0] = conv state: (conv_dim_local, conv_rows) in DS layout.
100 - shapes[1] = SSM temporal: (local_num_heads, head_dim).
115 + shapes[1] = temporal state (model-specific shape).
116 local_tp: this engine's tensor-parallel size.
117
118 Returns:
104 - MambaConvSplitInfo with per-rank x_local, b_local, conv_rows,
119 + MambaConvSplitInfo with per-rank sub-projection dims, conv_rows,
120 conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes).
121 """
107 - if mamba_spec.mamba_type != MambaAttentionBackendEnum.MAMBA2:
122 + _supported = (
123 + MambaAttentionBackendEnum.MAMBA2,
124 + MambaAttentionBackendEnum.GDN_ATTN,
125 + )
126 + if mamba_spec.mamba_type not in _supported:
127 raise NotImplementedError(
109 - f"3-read conv transfer only supports Mamba2 models, "
110 - f"got mamba_type={mamba_spec.mamba_type!r}. "
111 - f"Mamba1 SSM temporal shape is (intermediate_size // tp, state_size) "
112 - f"which cannot be used to reconstruct intermediate_size."
128 + f"3-read conv transfer only supports Mamba2 and GDN models, "
129 + f"got mamba_type={mamba_spec.mamba_type!r}."
130 )
131
132 conv_shape = mamba_spec.shapes[0]
@@ -121,29 +138,10 @@ def derive_mamba_conv_split(
138 local_conv_dim = conv_shape[0] # DS: (conv_dim_local, conv_rows)
139 conv_rows = conv_shape[1]
140
124 - # NOTE (ZhanqiuHu): intermediate_size (= global x dim) is not stored
125 - # in MambaSpec, so we reconstruct it from the SSM temporal state shape:
126 - # shapes[1] = (local_num_heads, head_dim), already divided by TP.
127 - head_dim = mamba_spec.shapes[1][1]
128 - local_num_heads = mamba_spec.shapes[1][0]
129 - intermediate_size = local_num_heads * local_tp * head_dim
130 -
131 - # NOTE (ZhanqiuHu): global conv dim = intermediate_size + 2 * groups_ss,
132 - # where groups_ss is the B (= C) dimension. B and C are always the same
133 - # size, so we recover groups_ss from the remainder after subtracting x.
134 - remainder = local_conv_dim * local_tp - intermediate_size
135 - assert remainder > 0 and remainder % 2 == 0, (
136 - f"Conv dim ({local_conv_dim}*tp={local_tp}) doesn't decompose into "
137 - f"intermediate_size={intermediate_size} + 2*groups_ss. "
138 - f"remainder={remainder}"
139 - )
140 - groups_ss = remainder // 2
141 -
141 conv_dtype_size = torch.tensor(
142 [],
143 dtype=mamba_spec.dtypes[0], # type: ignore[misc]
144 ).element_size()
146 -
145 ssm_dtype_size = torch.tensor(
146 [],
147 dtype=mamba_spec.dtypes[1], # type: ignore[misc]
@@ -151,11 +149,54 @@ def derive_mamba_conv_split(
149 conv_state_bytes = torch.Size(mamba_spec.shapes[0]).numel() * conv_dtype_size
150 ssm_state_bytes = torch.Size(mamba_spec.shapes[1]).numel() * ssm_dtype_size
151
154 - # Divide by TP to get per-rank column counts.
152 + if mamba_spec.mamba_type == MambaAttentionBackendEnum.MAMBA2:
153 + # NOTE (ZhanqiuHu): intermediate_size (= global x dim) is not stored
154 + # in MambaSpec, so we reconstruct it from the SSM temporal state shape:
155 + # shapes[1] = (local_num_heads, head_dim), already divided by TP.
156 + head_dim = mamba_spec.shapes[1][1]
157 + local_num_heads = mamba_spec.shapes[1][0]
158 + intermediate_size = local_num_heads * local_tp * head_dim
159 +
160 + # NOTE (ZhanqiuHu): global conv dim = intermediate_size + 2 * groups_ss,
161 + # where groups_ss is the B (= C) dimension. B and C are always the same
162 + # size, so we recover groups_ss from the remainder after subtracting x.
163 + remainder = local_conv_dim * local_tp - intermediate_size
164 + assert remainder > 0 and remainder % 2 == 0, (
165 + f"Conv dim ({local_conv_dim}*tp={local_tp}) doesn't decompose "
166 + f"into intermediate_size={intermediate_size} + 2*groups_ss. "
167 + f"remainder={remainder}"
168 + )
169 + groups_ss = remainder // 2
170 +
171 + # Divide by TP to get per-rank column counts.
172 + x_local = intermediate_size // local_tp
173 + b_local = groups_ss // local_tp
174 + local_proj_dims = (x_local, b_local, b_local)
175 + elif mamba_spec.mamba_type == MambaAttentionBackendEnum.GDN_ATTN:
176 + # GDN: conv = [Q, K, V] where dim(Q) == dim(K) == key_dim.
177 + # conv_dim = key_dim*2 + value_dim (all global, divided by TP).
178 + # Temporal state shape is (num_v_heads/TP, head_v_dim, head_k_dim).
179 + temporal_shape = mamba_spec.shapes[1]
180 + num_v_heads_local = temporal_shape[0]
181 + head_v_dim = temporal_shape[1]
182 + value_dim_local = num_v_heads_local * head_v_dim
183 +
184 + remainder = local_conv_dim - value_dim_local
185 + assert remainder > 0 and remainder % 2 == 0, (
186 + f"GDN conv dim ({local_conv_dim}) doesn't decompose into "
187 + f"2*key_dim_local + value_dim_local={value_dim_local}. "
188 + f"remainder={remainder}"
189 + )
190 + key_dim_local = remainder // 2
191 + local_proj_dims = (key_dim_local, key_dim_local, value_dim_local)
192 + else:
193 + raise NotImplementedError(
194 + f"Conv split not supported for mamba_type={mamba_spec.mamba_type!r}"
195 + )
196 +
197 return MambaConvSplitInfo(
198 conv_rows=conv_rows,
157 - x_local=intermediate_size // local_tp,
158 - b_local=groups_ss // local_tp,
199 + local_proj_dims=local_proj_dims,
200 conv_dtype_size=conv_dtype_size,
201 ssm_sizes=(conv_state_bytes, ssm_state_bytes),
202 )