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 5, 03:30 UTC
[Bugfix][NIXL] Fix `_logical_to_kernel_block_ids` conversion for non-mamba models (#39724)
Signed-off-by: Zhanqiu Hu <zhu@redhat.com>
zhanqiuhu committed
Apr 15, 2026 at 16:08 UTC
41488f2acdc53eadbae98a316df47ba039589fe7
2 files changed
+97
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+93
@@ -89,6 +89,99 @@ def test_logical_to_kernel_block_ids_with_hma():
89
)
90
91
92
+@pytest.mark.cpu_test
93
+@pytest.mark.parametrize(
94
+ "has_mamba,swa_enabled,mamba_enabled,remote_ratio,"
95
+ "remote_block_ids,expected_remote_block_ids",
96
+ [
97
+ # Non-mamba (FA+SWA): both groups expanded via _logical_to_kernel_block_ids.
98
+ # Regression for https://github.com/vllm-project/vllm/pull/39724
99
+ (
100
+ False,
101
+ True,
102
+ False,
103
+ 1,
104
+ ([0, 1, 2], [3, 4]),
105
+ [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]],
106
+ ),
107
+ # Mamba (FA+Mamba): FA expanded via _logical_to_remote_kernel_block_ids,
108
+ # Mamba passed through unchanged.
109
+ # remote_ratio=261 (Nemotron 30B TP=1) != local_ratio=2 so that using
110
+ # the wrong conversion method produces different FA results.
111
+ (
112
+ True,
113
+ False,
114
+ True,
115
+ 261,
116
+ ([0, 1, 2], [10, 11]),
117
+ [[0, 1, 261, 262, 522, 523], [10, 11]],
118
+ ),
119
+ ],
120
+ ids=["non_mamba_fa_swa", "mamba_fa_ssm"],
121
+)
122
+def test_read_blocks_for_req_expands_remote_ids(
123
+ has_mamba,
124
+ swa_enabled,
125
+ mamba_enabled,
126
+ remote_ratio,
127
+ remote_block_ids,
128
+ expected_remote_block_ids,
129
+):
130
+ """_read_blocks_for_req must expand remote logical block IDs to kernel
131
+ block IDs when kernel block size != logical block size.
132
+
133
+ Non-mamba path uses _logical_to_kernel_block_ids (all groups expanded).
134
+ Mamba path uses _logical_to_remote_kernel_block_ids (FA expanded, Mamba
135
+ passed through).
136
+ """
137
+ from unittest.mock import MagicMock
138
+
139
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
140
+ NixlConnectorMetadata,
141
+ )
142
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
143
+ NixlConnectorWorker,
144
+ )
145
+
146
+ worker = object.__new__(NixlConnectorWorker)
147
+ worker._has_mamba = has_mamba
148
+ worker._physical_blocks_per_logical_kv_block = 2
149
+ worker.kv_cache_config = make_kv_cache_config(
150
+ block_size=16, swa_enabled=swa_enabled, mamba_enabled=mamba_enabled
151
+ )
152
+
153
+ remote_engine_id = "remote-engine"
154
+ if has_mamba:
155
+ worker._mamba_phys_ratio = {remote_engine_id: remote_ratio}
156
+
157
+ # Mock kv_topo: empty remote ranks skips the transfer machinery entirely,
158
+ # isolating the block-ID expansion logic.
159
+ worker.kv_topo = MagicMock()
160
+ worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = []
161
+ worker.kv_topo.tp_ratio_from_engine_id.return_value = 1
162
+
163
+ metadata = NixlConnectorMetadata()
164
+ metadata.add_new_req_to_recv(
165
+ request_id="test-req",
166
+ local_block_ids=([0, 1], [2, 3]),
167
+ kv_transfer_params={
168
+ "remote_block_ids": remote_block_ids,
169
+ "remote_engine_id": remote_engine_id,
170
+ "remote_request_id": "prefill-test-req",
171
+ "remote_host": "localhost",
172
+ "remote_port": 1234,
173
+ "tp_size": 1,
174
+ },
175
+ )
176
+
177
+ meta = metadata.reqs_to_recv["test-req"]
178
+ worker._read_blocks_for_req("test-req", meta)
179
+
180
+ assert meta.remote.block_ids == expected_remote_block_ids, (
181
+ f"Expected {expected_remote_block_ids}, got {meta.remote.block_ids}"
182
+ )
183
+
184
+
185
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
186
def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
187
"""Test that a prefill instance returns fewer "remote blocks" for the SWA groups
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+4
@@ -1914,6 +1914,10 @@ class NixlConnectorWorker:
1914
meta.remote.block_ids,
1915
self._mamba_phys_ratio[meta.remote.engine_id],
1916
)
1917
+ else:
1918
+ meta.remote.block_ids = self._logical_to_kernel_block_ids(
1919
+ meta.remote.block_ids
1920
+ )
1921
# D may have to perform multiple reads from different remote ranks.
1922
for i, remote_rank in enumerate(remote_ranks):
1923
if self.use_mla and tp_ratio < 0 and i > 0: