PD disagg with Mooncake Connector: GDN support (Qwen3.5) and MLA support (Deepseek-V4-Flash) (#46807)
Dakai An committed
Jun 30, 2026 at 14:29 UTC
0feca7ffa8f626f51f5ea262eb586a3bbe704991
7 files changed
+735
-103
tests/v1/kv_connector/unit/test_mooncake_connector.py
+36
-4
@@ -33,14 +33,37 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import
33
)
34
from vllm.utils.network_utils import get_open_port
35
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
36
-from vllm.v1.kv_cache_interface import KVCacheConfig
36
+from vllm.v1.kv_cache_interface import (
37
+ FullAttentionSpec,
38
+ KVCacheConfig,
39
+ KVCacheGroupSpec,
40
+)
41
from vllm.v1.request import RequestStatus
42
43
from .utils import create_request, create_scheduler, create_vllm_config
44
45
46
def _make_test_kv_cache_config() -> KVCacheConfig:
43
- return KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[])
47
+ return KVCacheConfig(
48
+ num_blocks=0,
49
+ kv_cache_tensors=[],
50
+ kv_cache_groups=[
51
+ KVCacheGroupSpec(
52
+ [
53
+ "model.layers.0.self_attn",
54
+ "model.layers.1.self_attn",
55
+ "model.layers.0.mla_attn",
56
+ "model.layers.1.eagle_attn",
57
+ ],
58
+ FullAttentionSpec(
59
+ block_size=16,
60
+ num_kv_heads=4,
61
+ head_size=64,
62
+ dtype=torch.float16,
63
+ ),
64
+ )
65
+ ],
66
+ )
67
68
69
class FakeMooncakeWrapper:
@@ -126,6 +149,8 @@ async def test_build_transfer_params_separates_prefill_pp_layers():
149
worker.is_kv_producer = True
150
worker.tp_rank = 0
151
worker.tp_size = 1
152
+ worker.kv_cache_config = _make_test_kv_cache_config()
153
+ worker._physical_blocks_per_logical_kv_block = 1
154
worker.transfer_topo = SimpleNamespace(local_replicates_kv_cache=False)
155
156
block_len = 256
@@ -206,6 +231,7 @@ async def test_build_transfer_params_separates_prefill_pp_layers():
231
req_blocks={"d-req-pp": (transfer_id, [[20, 21]])},
232
kv_caches_base_addr=[region.base_addr for region in remote_regions],
233
block_lens=[region.block_len for region in remote_regions],
234
+ kv_block_lens=[region.kv_block_len for region in remote_regions],
235
registered_layer_names=[region.layer_name for region in remote_regions],
236
registered_layer_indices=[region.layer_index for region in remote_regions],
237
)
@@ -266,6 +292,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
292
kv_half = block_len // 2
293
prefill_worker.kv_caches_base_addr = [0x1000]
294
prefill_worker.block_len_per_layer = [block_len]
295
+ prefill_worker.kv_block_len_per_layer = [kv_half]
296
prefill_worker.registered_layer_names = ["model.layers.1.self_attn"]
297
prefill_worker.registered_layer_indices = [1]
298
@@ -294,6 +321,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
321
req_blocks={"d-req-layer-align": (transfer_id, [[20]])},
322
kv_caches_base_addr=[0xA000, 0xB000],
323
block_lens=[block_len, block_len],
324
+ kv_block_lens=[kv_half, kv_half],
325
registered_layer_names=[
326
"model.layers.0.self_attn",
327
"model.layers.1.self_attn",
@@ -804,7 +832,9 @@ async def test_kv_producer(monkeypatch):
832
prefill_worker = prefill_connector.connector_worker
833
prefill_worker.kv_caches_base_addr = [0x1000]
834
block_len = 4096
835
+ kv_half = block_len // 2
836
prefill_worker.block_len_per_layer = [block_len]
837
+ prefill_worker.kv_block_len_per_layer = [kv_half]
838
prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
839
prefill_worker.registered_layer_indices = [0]
840
@@ -832,6 +862,7 @@ async def test_kv_producer(monkeypatch):
862
req_blocks={"d-req-1": (transfer_id, [[20, 21]])},
863
kv_caches_base_addr=[0x2000],
864
block_lens=[block_len],
865
+ kv_block_lens=[kv_half],
866
registered_layer_names=["model.layers.0.self_attn"],
867
registered_layer_indices=[0],
868
)
@@ -845,8 +876,6 @@ async def test_kv_producer(monkeypatch):
876
) as mock_send_blocks:
877
# With blocks-first layout, each block is virtually split
878
# into K and V halves, producing non-coalesced transfers.
848
- kv_half = block_len // 2
849
-
879
def expected_split_transfers(src_base, dst_base, src_blocks, dst_blocks):
880
"""Build expected (src_ptrs, dst_ptrs, lengths) for
881
virtual-split K/V transfers."""
@@ -981,6 +1010,7 @@ async def test_kv_consumuer(monkeypatch):
1010
decode_worker = decode_connector.connector_worker
1011
decode_worker.kv_caches_base_addr = [0x1000]
1012
decode_worker.block_len_per_layer = [4096]
1013
+ decode_worker.kv_block_len_per_layer = [4096]
1014
decode_worker.registered_layer_names = ["model.layers.0.self_attn"]
1015
decode_worker.registered_layer_indices = [0]
1016
decode_worker.rpc_port = 54321
@@ -1236,6 +1266,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
1266
1267
prefill_worker.kv_caches_base_addr = [0x1000]
1268
prefill_worker.block_len_per_layer = [local_block_len]
1269
+ prefill_worker.kv_block_len_per_layer = [local_block_len // 2]
1270
prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
1271
prefill_worker.registered_layer_indices = [0]
1272
@@ -1283,6 +1314,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
1314
},
1315
kv_caches_base_addr=[0x2000],
1316
block_lens=[remote_block_len],
1317
+ kv_block_lens=[remote_block_len // 2],
1318
registered_layer_names=["model.layers.0.self_attn"],
1319
registered_layer_indices=[0],
1320
)
tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
+2
@@ -257,6 +257,7 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch):
257
},
258
kv_caches_base_addr=[0x2000],
259
block_lens=[block_len],
260
+ kv_block_lens=[block_len],
261
)
262
263
local_regions = [
@@ -348,6 +349,7 @@ async def test_build_transfer_params_group_count_mismatch(monkeypatch):
349
},
350
kv_caches_base_addr=[0x2000],
351
block_lens=[block_len],
352
+ kv_block_lens=[block_len],
353
)
354
355
local_regions = [
tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
new
+386
@@ -0,0 +1,386 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for MooncakeConnector hybrid FA + GDN support.
4
+
5
+GDN is represented as a MambaSpec in vLLM, so these tests exercise the
6
+Mooncake MambaSpec path with mamba_type=GDN_ATTN. Mamba2 is intentionally not
7
+validated by this test module.
8
+"""
9
+
10
+import asyncio
11
+from types import SimpleNamespace
12
+from unittest.mock import patch
13
+
14
+import pytest
15
+import torch
16
+
17
+from vllm.config import set_current_vllm_config
18
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
19
+ KVConnectorRole,
20
+ MooncakeConnector,
21
+ MooncakeConnectorScheduler,
22
+ MooncakeConnectorWorker,
23
+ MooncakeXferMetadata,
24
+ SendBlockMeta,
25
+ TransferRegion,
26
+)
27
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
28
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
29
+from vllm.v1.kv_cache_interface import (
30
+ FullAttentionSpec,
31
+ KVCacheConfig,
32
+ KVCacheGroupSpec,
33
+ MambaSpec,
34
+)
35
+
36
+from .test_mooncake_connector import patch_worker_dependencies
37
+from .utils import create_request, create_vllm_config
38
+
39
+
40
+def noop_shutdown():
41
+ pass
42
+
43
+
44
+def make_hybrid_gdn_kv_cache_config(block_size: int) -> KVCacheConfig:
45
+ gdn_spec = MambaSpec(
46
+ block_size=block_size,
47
+ shapes=((6, 3), (1, 2, 2)),
48
+ dtypes=(torch.float16, torch.float16),
49
+ mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
50
+ )
51
+ assert gdn_spec.mamba_type == MambaAttentionBackendEnum.GDN_ATTN
52
+ return KVCacheConfig(
53
+ num_blocks=16,
54
+ kv_cache_tensors=[],
55
+ kv_cache_groups=[
56
+ KVCacheGroupSpec(
57
+ ["model.layers.0.self_attn"],
58
+ FullAttentionSpec(
59
+ block_size=block_size,
60
+ num_kv_heads=1,
61
+ head_size=1,
62
+ dtype=torch.float16,
63
+ ),
64
+ ),
65
+ KVCacheGroupSpec(
66
+ ["model.layers.1.linear_attn"],
67
+ gdn_spec,
68
+ ),
69
+ ],
70
+ )
71
+
72
+
73
+def make_hybrid_gdn_scheduler(kv_role: str) -> MooncakeConnectorScheduler:
74
+ vllm_config = create_vllm_config(
75
+ kv_connector="MooncakeConnector",
76
+ kv_role=kv_role,
77
+ )
78
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
79
+ return MooncakeConnectorScheduler(
80
+ vllm_config=vllm_config,
81
+ engine_id="test-engine",
82
+ kv_cache_config=make_hybrid_gdn_kv_cache_config(
83
+ vllm_config.cache_config.block_size
84
+ ),
85
+ )
86
+
87
+
88
+@pytest.mark.cpu_test
89
+def test_hybrid_gdn_remote_prefill_uses_mamba_n_minus_one():
90
+ scheduler = make_hybrid_gdn_scheduler(kv_role="kv_consumer")
91
+ request = create_request(num_tokens=10, do_remote_prefill=True)
92
+
93
+ num_new_tokens, is_async = scheduler.get_num_new_matched_tokens(
94
+ request, num_computed_tokens=0
95
+ )
96
+
97
+ assert num_new_tokens == request.num_prompt_tokens - 1
98
+ assert is_async is True
99
+
100
+
101
+@pytest.mark.cpu_test
102
+def test_hybrid_gdn_remote_decode_truncates_prefill_once():
103
+ scheduler = make_hybrid_gdn_scheduler(kv_role="kv_producer")
104
+ request = create_request(num_tokens=10, do_remote_decode=True)
105
+ original_tokens = list(request.prompt_token_ids)
106
+
107
+ num_new_tokens, is_async = scheduler.get_num_new_matched_tokens(
108
+ request, num_computed_tokens=0
109
+ )
110
+
111
+ assert num_new_tokens == 0
112
+ assert is_async is False
113
+ assert request.prompt_token_ids == original_tokens[:-1]
114
+ assert request._all_token_ids == original_tokens[:-1]
115
+ assert request.num_prompt_tokens == len(original_tokens) - 1
116
+ assert request.max_tokens == 1
117
+ assert request.kv_transfer_params["_p_side_truncated"] is True
118
+
119
+ scheduler.get_num_new_matched_tokens(request, num_computed_tokens=0)
120
+ assert request.prompt_token_ids == original_tokens[:-1]
121
+
122
+
123
+def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
124
+ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
125
+ vllm_config = create_vllm_config(
126
+ kv_connector="MooncakeConnector",
127
+ kv_role="kv_consumer",
128
+ )
129
+ kv_cache_config = make_hybrid_gdn_kv_cache_config(
130
+ vllm_config.cache_config.block_size
131
+ )
132
+
133
+ with set_current_vllm_config(vllm_config), patch_worker_dependencies():
134
+ connector = MooncakeConnector(
135
+ vllm_config,
136
+ KVConnectorRole.WORKER,
137
+ kv_cache_config,
138
+ )
139
+ worker = connector.connector_worker
140
+
141
+ fa_cache = torch.empty((2, 2, 11), dtype=torch.float16)
142
+ gdn_conv_state = torch.empty((2, 22), dtype=torch.float16)
143
+ gdn_ssm_state = torch.empty((2, 4), dtype=torch.float16)
144
+
145
+ worker.register_kv_caches(
146
+ {
147
+ "model.layers.0.self_attn": fa_cache,
148
+ "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
149
+ }
150
+ )
151
+
152
+ assert worker.transfer_topo.is_mamba is True
153
+ assert worker.registered_layer_names == [
154
+ "model.layers.0.self_attn",
155
+ "model.layers.1.linear_attn",
156
+ ]
157
+ assert worker.registered_group_indices == [0, 1]
158
+ assert worker.kv_caches_base_addr == [
159
+ fa_cache.data_ptr(),
160
+ gdn_conv_state.data_ptr(),
161
+ ]
162
+
163
+ worker.shutdown()
164
+ worker.shutdown = noop_shutdown
165
+ connector.connector_worker = None
166
+
167
+
168
+def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
169
+ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
170
+ vllm_config = create_vllm_config(
171
+ kv_connector="MooncakeConnector",
172
+ kv_role="kv_consumer",
173
+ )
174
+ kv_cache_config = make_hybrid_gdn_kv_cache_config(
175
+ vllm_config.cache_config.block_size
176
+ )
177
+
178
+ with set_current_vllm_config(vllm_config), patch_worker_dependencies():
179
+ connector = MooncakeConnector(
180
+ vllm_config,
181
+ KVConnectorRole.WORKER,
182
+ kv_cache_config,
183
+ )
184
+ worker = connector.connector_worker
185
+
186
+ backing = torch.empty((4, 64), dtype=torch.float16)
187
+ fa_cache = backing[:2, :16]
188
+ gdn_conv_state = backing[:3]
189
+ gdn_ssm_state = torch.empty((3, 4), dtype=torch.float16)
190
+
191
+ with patch.object(
192
+ worker.engine, "batch_register_memory", return_value=0
193
+ ) as batch_register_memory:
194
+ worker.register_kv_caches(
195
+ {
196
+ "model.layers.0.self_attn": fa_cache,
197
+ "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
198
+ }
199
+ )
200
+
201
+ assert worker.kv_caches_base_addr == [
202
+ fa_cache.data_ptr(),
203
+ gdn_conv_state.data_ptr(),
204
+ ]
205
+ batch_register_memory.assert_called_once()
206
+ registered_ptrs, registered_lens = batch_register_memory.call_args[0]
207
+ assert registered_ptrs == [backing.data_ptr()]
208
+ assert registered_lens == [backing.untyped_storage().nbytes()]
209
+
210
+ worker.shutdown()
211
+ worker.shutdown = noop_shutdown
212
+ connector.connector_worker = None
213
+
214
+
215
+def test_hybrid_gdn_transfer_params_preserve_group_identity(monkeypatch):
216
+ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
217
+ vllm_config = create_vllm_config(
218
+ kv_connector="MooncakeConnector",
219
+ kv_role="kv_producer",
220
+ )
221
+ kv_cache_config = make_hybrid_gdn_kv_cache_config(
222
+ vllm_config.cache_config.block_size
223
+ )
224
+
225
+ with set_current_vllm_config(vllm_config), patch_worker_dependencies():
226
+ connector = MooncakeConnector(
227
+ vllm_config,
228
+ KVConnectorRole.WORKER,
229
+ kv_cache_config,
230
+ )
231
+ worker = connector.connector_worker
232
+
233
+ block_len = 0x100
234
+ transfer_id = "xfer-hybrid-gdn"
235
+
236
+ async def build_transfer_params():
237
+ send_meta = SendBlockMeta(
238
+ p_req_id="p-hybrid-gdn",
239
+ transfer_id=transfer_id,
240
+ local_block_ids=[
241
+ [10, 11],
242
+ [NULL_BLOCK_ID, 4],
243
+ ],
244
+ ready=asyncio.Event(),
245
+ )
246
+ return await worker._build_transfer_params(
247
+ [("d-hybrid-gdn", send_meta)],
248
+ xfer_meta,
249
+ local_regions,
250
+ remote_regions,
251
+ )
252
+
253
+ xfer_meta = MooncakeXferMetadata(
254
+ remote_hostname="consumer-host",
255
+ remote_port=54321,
256
+ remote_tp_size=1,
257
+ remote_tp_rank=0,
258
+ req_blocks={
259
+ "d-hybrid-gdn": (
260
+ transfer_id,
261
+ [
262
+ [30, 31],
263
+ [NULL_BLOCK_ID, 7],
264
+ ],
265
+ )
266
+ },
267
+ kv_caches_base_addr=[],
268
+ block_lens=[],
269
+ kv_block_lens=[],
270
+ )
271
+
272
+ local_regions = [
273
+ TransferRegion(
274
+ layer_name="model.layers.1.linear_attn",
275
+ layer_index=1,
276
+ base_addr=0x5000,
277
+ block_len=block_len,
278
+ kv_block_len=block_len,
279
+ group_index=1,
280
+ ),
281
+ TransferRegion(
282
+ layer_name="model.layers.0.self_attn",
283
+ layer_index=0,
284
+ base_addr=0x1000,
285
+ block_len=block_len,
286
+ kv_block_len=block_len,
287
+ group_index=0,
288
+ ),
289
+ ]
290
+ remote_regions = [
291
+ TransferRegion(
292
+ layer_name="model.layers.1.linear_attn",
293
+ layer_index=1,
294
+ base_addr=0x6000,
295
+ block_len=block_len,
296
+ kv_block_len=block_len,
297
+ group_index=1,
298
+ ),
299
+ TransferRegion(
300
+ layer_name="model.layers.0.self_attn",
301
+ layer_index=0,
302
+ base_addr=0x2000,
303
+ block_len=block_len,
304
+ kv_block_len=block_len,
305
+ group_index=0,
306
+ ),
307
+ ]
308
+
309
+ src_ptrs, dst_ptrs, lengths, err_reqs, err_msg = asyncio.run(
310
+ build_transfer_params()
311
+ )
312
+
313
+ assert err_reqs == []
314
+ assert err_msg is None
315
+ assert src_ptrs == [
316
+ 0x5000 + 4 * block_len,
317
+ 0x1000 + 10 * block_len,
318
+ ]
319
+ assert dst_ptrs == [
320
+ 0x6000 + 7 * block_len,
321
+ 0x2000 + 30 * block_len,
322
+ ]
323
+ assert lengths == [block_len, 2 * block_len]
324
+
325
+ worker.shutdown()
326
+ worker.shutdown = noop_shutdown
327
+ connector.connector_worker = None
328
+
329
+
330
+def test_logical_to_kernel_block_ids_expands_fa_not_gdn():
331
+ worker = object.__new__(MooncakeConnectorWorker)
332
+ worker.shutdown = noop_shutdown
333
+ worker._physical_blocks_per_logical_kv_block = 17
334
+ worker.kv_cache_config = make_hybrid_gdn_kv_cache_config(block_size=544)
335
+
336
+ block_ids = [[2], [2]]
337
+ kernel_block_ids = worker._logical_to_kernel_block_ids(block_ids)
338
+
339
+ assert kernel_block_ids == [list(range(34, 51)), [2]]
340
+
341
+
342
+def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
343
+ monkeypatch,
344
+):
345
+ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
346
+ vllm_config = create_vllm_config(
347
+ kv_connector="MooncakeConnector",
348
+ kv_role="kv_producer",
349
+ )
350
+ kv_cache_config = make_hybrid_gdn_kv_cache_config(
351
+ vllm_config.cache_config.block_size
352
+ )
353
+
354
+ with set_current_vllm_config(vllm_config), patch_worker_dependencies():
355
+ connector = MooncakeConnector(
356
+ vllm_config,
357
+ KVConnectorRole.WORKER,
358
+ kv_cache_config,
359
+ )
360
+ worker = connector.connector_worker
361
+
362
+ worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=True)
363
+ regions = worker._get_transfer_regions(
364
+ base_addrs=[0x1000, 0x2000],
365
+ block_lens=[0x100, 0x100],
366
+ kv_block_lens=[0x40, 0x100],
367
+ layer_names=[
368
+ "model.layers.0.self_attn",
369
+ "model.layers.1.linear_attn",
370
+ ],
371
+ layer_indices=[0, 1],
372
+ group_indices=[0, 1],
373
+ )
374
+
375
+ assert [
376
+ (region.group_index, region.base_addr, region.kv_block_len)
377
+ for region in regions
378
+ ] == [
379
+ (0, 0x1000, 0x40),
380
+ (0, 0x1040, 0x40),
381
+ (1, 0x2000, 0x100),
382
+ ]
383
+
384
+ worker.shutdown()
385
+ worker.shutdown = noop_shutdown
386
+ connector.connector_worker = None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+308
-96
@@ -22,7 +22,6 @@ from vllm.config import VllmConfig
22
from vllm.distributed.kv_transfer.kv_connector.utils import (
23
EngineId,
24
TransferTopology,
25
- get_current_attn_backend,
25
get_current_attn_backends,
26
)
27
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -51,10 +50,18 @@ from vllm.platforms import current_platform
50
from vllm.utils.math_utils import cdiv
51
from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
52
from vllm.v1.attention.backend import AttentionMetadata
54
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
53
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, get_kv_cache_layout
54
from vllm.v1.core.sched.output import SchedulerOutput
56
-from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
55
+from vllm.v1.kv_cache_interface import (
56
+ FullAttentionSpec,
57
+ KVCacheSpec,
58
+ MambaSpec,
59
+ MLAAttentionSpec,
60
+ SlidingWindowMLASpec,
61
+ SlidingWindowSpec,
62
+)
63
from vllm.v1.request import RequestStatus
64
+from vllm.v1.worker.block_table import BlockTable
65
from vllm.v1.worker.utils import select_common_block_size
66
67
logger = init_logger(__name__)
@@ -85,6 +92,7 @@ class TransferRegion:
92
base_addr: int
93
block_len: int
94
kv_block_len: int
95
+ group_index: int = 0
96
97
98
def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int:
@@ -111,24 +119,58 @@ def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int:
119
def _expand_transfer_regions(
120
base_addrs: list[int],
121
block_lens: list[int],
122
+ kv_block_lens: list[int],
123
layer_names: list[str],
124
layer_indices: list[int],
125
is_kv_layout_blocks_first: bool,
126
+ group_indices: list[int] | None = None,
127
+ split_kv_regions: list[bool] | None = None,
128
) -> list[TransferRegion]:
129
"""Expand registered KV tensors into the regions transferred by Mooncake."""
130
assert (
120
- len(base_addrs) == len(block_lens) == len(layer_names) == len(layer_indices)
131
+ len(base_addrs)
132
+ == len(block_lens)
133
+ == len(kv_block_lens)
134
+ == len(layer_names)
135
+ == len(layer_indices)
136
), (
137
"Mooncake transfer regions require matching metadata lengths, got "
138
f"base_addrs={len(base_addrs)}, block_lens={len(block_lens)}, "
139
+ f"kv_block_lens={len(kv_block_lens)}, "
140
f"layer_names={len(layer_names)}, "
141
f"layer_indices={len(layer_indices)}."
142
)
143
+ if group_indices is None:
144
+ group_indices = [0] * len(layer_names)
145
+ assert len(group_indices) == len(layer_names), (
146
+ "Mooncake transfer regions require matching group metadata lengths, "
147
+ f"got group_indices={len(group_indices)}, layer_names={len(layer_names)}."
148
+ )
149
+ if split_kv_regions is None:
150
+ split_kv_regions = [is_kv_layout_blocks_first] * len(layer_names)
151
+ assert len(split_kv_regions) == len(layer_names), (
152
+ "Mooncake transfer regions require matching split metadata, "
153
+ f"got split_kv_regions={len(split_kv_regions)}, "
154
+ f"layer_names={len(layer_names)}."
155
+ )
156
regions: list[TransferRegion] = []
128
- for base_addr, block_len, layer_name, layer_index in zip(
129
- base_addrs, block_lens, layer_names, layer_indices
157
+ for (
158
+ base_addr,
159
+ block_len,
160
+ kv_block_len,
161
+ layer_name,
162
+ layer_index,
163
+ group_index,
164
+ split_kv_region,
165
+ ) in zip(
166
+ base_addrs,
167
+ block_lens,
168
+ kv_block_lens,
169
+ layer_names,
170
+ layer_indices,
171
+ group_indices,
172
+ split_kv_regions,
173
):
131
- kv_block_len = block_len // 2 if is_kv_layout_blocks_first else block_len
174
regions.append(
175
TransferRegion(
176
layer_name=layer_name,
@@ -136,9 +178,10 @@ def _expand_transfer_regions(
178
base_addr=base_addr,
179
block_len=block_len,
180
kv_block_len=kv_block_len,
181
+ group_index=group_index,
182
)
183
)
141
- if is_kv_layout_blocks_first:
184
+ if split_kv_region:
185
regions.append(
186
TransferRegion(
187
layer_name=layer_name,
@@ -146,6 +189,7 @@ def _expand_transfer_regions(
189
base_addr=base_addr + kv_block_len,
190
block_len=block_len,
191
kv_block_len=kv_block_len,
192
+ group_index=group_index,
193
)
194
)
195
return regions
@@ -308,6 +352,17 @@ def _align_transfer_regions(
352
f"{remote_region.layer_index}."
353
),
354
)
355
+ if local_region.group_index != remote_region.group_index:
356
+ return (
357
+ [],
358
+ [],
359
+ (
360
+ "Mooncake registered group index mismatch for "
361
+ f"{local_region.layer_name}: producer="
362
+ f"{local_region.group_index}, consumer="
363
+ f"{remote_region.group_index}."
364
+ ),
365
+ )
366
aligned_local.append(local_region)
367
aligned_remote.append(remote_region)
368
@@ -332,8 +387,10 @@ class MooncakeXferMetadata(
387
req_blocks: dict[ReqId, tuple[TransferId, list[list[int]]]]
388
kv_caches_base_addr: list[int]
389
block_lens: list[int]
390
+ kv_block_lens: list[int]
391
registered_layer_names: list[str] = msgspec.field(default_factory=list)
392
registered_layer_indices: list[int] = msgspec.field(default_factory=list)
393
+ registered_group_indices: list[int] = msgspec.field(default_factory=list)
394
395
396
class MooncakeXferResponseStatus(IntEnum):
@@ -581,6 +638,9 @@ class MooncakeConnectorScheduler:
638
for g in kv_cache_config.kv_cache_groups
639
)
640
)
641
+ # GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec
642
+ # path is currently tested with GDN; Mamba2 is not validated yet.
643
+ self._has_mamba = kv_cache_config.has_mamba_layers
644
645
# Requests that need to start recv/send.
646
# New requests are added by update_state_after_alloc in
@@ -617,6 +677,38 @@ class MooncakeConnectorScheduler:
677
for i, blocks in enumerate(block_ids)
678
]
679
680
+ def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
681
+ """D-side only. Returns N-1 for Mamba models since the decoder
682
+ always recomputes the last token and must start from h(N-1)."""
683
+ if self._has_mamba and num_prompt_tokens > 1:
684
+ return num_prompt_tokens - 1
685
+ return num_prompt_tokens
686
+
687
+ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
688
+ """P-side only: drop the last prompt token so the prefiller computes
689
+ h(N-1) instead of h(N). The decoder recomputes the last token to
690
+ derive h(N) correctly.
691
+
692
+ Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
693
+ request is preempted and rescheduled."""
694
+ params = request.kv_transfer_params
695
+ if (
696
+ params is not None
697
+ and not params.get("_p_side_truncated")
698
+ and request.num_prompt_tokens > 1
699
+ ):
700
+ if request.prompt_token_ids is not None:
701
+ request.prompt_token_ids.pop()
702
+ elif request.prompt_embeds is not None:
703
+ request.prompt_embeds = request.prompt_embeds[:-1]
704
+ else:
705
+ return
706
+
707
+ request._all_token_ids.pop()
708
+ request.num_prompt_tokens -= 1
709
+ request.max_tokens = 1
710
+ params["_p_side_truncated"] = True
711
+
712
def get_num_new_matched_tokens(
713
self, request: "Request", num_computed_tokens: int
714
) -> tuple[int, bool]:
@@ -650,10 +742,15 @@ class MooncakeConnectorScheduler:
742
# Remote prefill: get all prompt blocks from remote.
743
assert not self.is_kv_producer
744
token_ids = request.prompt_token_ids or []
653
- count = len(token_ids) - num_computed_tokens
745
+ count = self._get_remote_prefill_token_count(len(token_ids)) - (
746
+ num_computed_tokens
747
+ )
748
if count > 0:
749
return count, True
750
751
+ if params.get("do_remote_decode") and self._has_mamba:
752
+ self._truncate_mamba_request_for_prefill(request)
753
+
754
# No remote prefill for this request.
755
return 0, False
756
@@ -802,7 +899,7 @@ class MooncakeConnectorWorker:
899
self,
900
vllm_config: VllmConfig,
901
engine_id: str,
805
- kv_cache_config: "KVCacheConfig | None" = None,
902
+ kv_cache_config: "KVCacheConfig",
903
):
904
if TransferEngine is None:
905
logger.error("Mooncake is not available")
@@ -831,10 +928,15 @@ class MooncakeConnectorWorker:
928
protocol = kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr]
929
"mooncake_protocol", "rdma"
930
)
931
+ device_name = kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr]
932
+ "device_name", ""
933
+ )
934
logger.info(
935
"The Mooncake Transfer Engine is using %s as its protocol.", protocol
936
)
837
- ret_value = self.engine.initialize(self.hostname, "P2PHANDSHAKE", protocol, "")
937
+ ret_value = self.engine.initialize(
938
+ self.hostname, "P2PHANDSHAKE", protocol, device_name
939
+ )
940
if ret_value != 0:
941
raise RuntimeError("Mooncake Transfer Engine initialization failed.")
942
@@ -852,10 +954,11 @@ class MooncakeConnectorWorker:
954
self.engine_id: EngineId = engine_id
955
self.tp_rank = get_tensor_model_parallel_rank()
956
self.tp_size = get_tensor_model_parallel_world_size()
855
- self.num_blocks = 0
957
self.block_len_per_layer: list[int] = []
958
+ self.kv_block_len_per_layer: list[int] = []
959
self.registered_layer_names: list[str] = []
960
self.registered_layer_indices: list[int] = []
961
+ self.registered_group_indices: list[int] = []
962
self.seen_base_addresses: list[int] = []
963
964
assert (parallel_config := vllm_config.parallel_config)
@@ -916,26 +1019,40 @@ class MooncakeConnectorWorker:
1019
self.cache_config = vllm_config.cache_config
1020
self.kv_cache_config = kv_cache_config
1021
self.use_mla = self.model_config.use_mla
1022
+ self._physical_blocks_per_logical_kv_block = 1
1023
self._sync_block_size_with_kernel()
1024
921
- # Get the attention backend from the first layer
922
- # NOTE (NickLucche) models with multiple backends are not supported yet
923
- backend = get_current_attn_backend(vllm_config)
924
- self.backend_name = backend.get_name()
1025
+ self.attn_backends = get_current_attn_backends(vllm_config)
1026
self.kv_cache_layout = get_kv_cache_layout()
926
- logger.debug("Detected attention backend %s", self.backend_name)
1027
+ logger.debug(
1028
+ "Detected attention backends %s",
1029
+ [backend.get_name() for backend in self.attn_backends],
1030
+ )
1031
logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
1032
1033
self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
1034
+ self._layer_specs: dict[str, KVCacheSpec] = {}
1035
+ for group in kv_cache_config.kv_cache_groups:
1036
+ group_spec = group.kv_cache_spec
1037
+ specs_by_layer = getattr(group_spec, "kv_cache_specs", {})
1038
+ for layer_name in group.layer_names:
1039
+ self._layer_specs[layer_name] = specs_by_layer.get(
1040
+ layer_name, group_spec
1041
+ )
1042
+ self._layer_group_indices: dict[str, int] = {
1043
+ layer: group_index
1044
+ for group_index, group in enumerate(kv_cache_config.kv_cache_groups)
1045
+ for layer in group.layer_names
1046
+ }
1047
self.transfer_topo = TransferTopology(
1048
tp_rank=self.tp_rank,
1049
tp_size=self.tp_size,
1050
block_size=self.block_size,
1051
engine_id=self.engine_id,
1052
is_mla=self.use_mla,
936
- is_mamba=False,
1053
+ is_mamba=kv_cache_config.has_mamba_layers,
1054
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
938
- attn_backends=[backend],
1055
+ attn_backends=self.attn_backends,
1056
)
1057
1058
self.async_zmq_ctx = zmq.asyncio.Context()
@@ -958,6 +1075,9 @@ class MooncakeConnectorWorker:
1075
kernel_block_size,
1076
)
1077
assert self.block_size > kernel_block_size
1078
+ self._physical_blocks_per_logical_kv_block = (
1079
+ self.block_size // kernel_block_size
1080
+ )
1081
self.block_size = kernel_block_size
1082
1083
def __del__(self):
@@ -1092,14 +1212,18 @@ class MooncakeConnectorWorker:
1212
local_regions = self._get_transfer_regions(
1213
self.kv_caches_base_addr,
1214
self.block_len_per_layer,
1215
+ self.kv_block_len_per_layer,
1216
self.registered_layer_names,
1217
self.registered_layer_indices,
1218
+ self.registered_group_indices,
1219
)
1220
remote_regions = self._get_transfer_regions(
1221
meta.kv_caches_base_addr,
1222
meta.block_lens,
1223
+ meta.kv_block_lens,
1224
meta.registered_layer_names,
1225
meta.registered_layer_indices,
1226
+ meta.registered_group_indices,
1227
)
1228
local_regions, remote_regions, align_err = _align_transfer_regions(
1229
local_regions, remote_regions
@@ -1271,6 +1395,32 @@ class MooncakeConnectorWorker:
1395
remote_tp_ranks,
1396
)
1397
1398
+ def _logical_to_kernel_block_ids(
1399
+ self, block_ids: list[list[int]]
1400
+ ) -> list[list[int]]:
1401
+ # For example, if a 544-token logical block is served by 32-token
1402
+ # FA kernel blocks, FA block id k expands to [17k, ..., 17k + 16],
1403
+ # while the matching Mamba/GDN state block remains k. Only attention
1404
+ # groups need logical block ids expanded to kernel block ids; Mamba/GDN
1405
+ # state block ids stay in the logical/page-id space.
1406
+ if self._physical_blocks_per_logical_kv_block == 1:
1407
+ return block_ids
1408
+
1409
+ block_arange = np.arange(self._physical_blocks_per_logical_kv_block).reshape(
1410
+ 1, -1
1411
+ )
1412
+ group_specs = self.kv_cache_config.kv_cache_groups
1413
+ return [
1414
+ BlockTable.map_to_kernel_blocks(
1415
+ np.array(group),
1416
+ self._physical_blocks_per_logical_kv_block,
1417
+ block_arange,
1418
+ ).tolist()
1419
+ if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
1420
+ else group
1421
+ for i, group in enumerate(block_ids)
1422
+ ]
1423
+
1424
async def _build_transfer_params(
1425
self,
1426
ready_reqs: list[tuple[ReqId, SendBlockMeta]],
@@ -1293,14 +1443,6 @@ class MooncakeConnectorWorker:
1443
):
1444
continue
1445
1296
- # Per-group partial hit trimming, then flatten.
1297
- # With HMA, groups share the same KV tensor but use different
1298
- # block ranges. We trim and concatenate so the coalescer and
1299
- # address math see one flat block list — same as non-HMA, but
1300
- # now including blocks from every group.
1301
- local_block_ids: list[int] = []
1302
- remote_block_ids: list[int] = []
1303
- has_block_error = False
1446
if len(send_meta.local_block_ids) != len(remote_block_ids_per_group):
1447
logger.error(
1448
"req %s: KV group count mismatch: local=%d, remote=%d",
@@ -1312,26 +1454,55 @@ class MooncakeConnectorWorker:
1454
if err_msg is None:
1455
err_msg = "KV group count mismatch"
1456
continue
1315
- for local_group, remote_group in zip(
1316
- send_meta.local_block_ids, remote_block_ids_per_group
1457
+
1458
+ # Keep KV-cache group identity. Hybrid/HMA groups can carry
1459
+ # different semantics (e.g. full-attention KV pages vs GDN/Mamba
1460
+ # inner-state slots), so their block IDs must not be flattened and
1461
+ # reused for every registered region.
1462
+ local_block_ids_by_group: list[list[int]] = []
1463
+ remote_block_ids_by_group: list[list[int]] = []
1464
+ has_block_error = False
1465
+ group_specs = self.kv_cache_config.kv_cache_groups
1466
+ for group_index, (local_group, remote_group) in enumerate(
1467
+ zip(send_meta.local_block_ids, remote_block_ids_per_group)
1468
):
1469
+ is_mamba_group = isinstance(
1470
+ group_specs[group_index].kv_cache_spec,
1471
+ MambaSpec,
1472
+ )
1473
+ if is_mamba_group:
1474
+ # Mamba/GDN prefix caching can use null blocks only as
1475
+ # align-mode placeholders. They do not carry transferable
1476
+ # state, so skip them on both producer and consumer sides.
1477
+ local_group = [
1478
+ block_id
1479
+ for block_id in local_group
1480
+ if block_id != NULL_BLOCK_ID
1481
+ ]
1482
+ remote_group = [
1483
+ block_id
1484
+ for block_id in remote_group
1485
+ if block_id != NULL_BLOCK_ID
1486
+ ]
1487
+
1488
n_local = len(local_group)
1489
n_remote = len(remote_group)
1490
if n_local < n_remote:
1491
logger.error(
1492
"req %s: local blocks(%d) < remote blocks(%d) "
1323
- "in a KV cache group",
1493
+ "in a KV cache group (is_mamba_group=%s)",
1494
d_req_id,
1495
n_local,
1496
n_remote,
1497
+ is_mamba_group,
1498
)
1499
has_block_error = True
1500
break
1330
- if n_local > n_remote:
1501
+ elif n_local > n_remote:
1502
# Partial prefix cache hit: just read uncomputed blocks.
1332
- local_group = local_group[-n_remote:]
1333
- local_block_ids.extend(local_group)
1334
- remote_block_ids.extend(remote_group)
1503
+ local_group = local_group[-n_remote:] if n_remote > 0 else []
1504
+ local_block_ids_by_group.append(local_group)
1505
+ remote_block_ids_by_group.append(remote_group)
1506
1507
if has_block_error:
1508
err_reqs.append(d_req_id)
@@ -1339,22 +1510,44 @@ class MooncakeConnectorWorker:
1510
err_msg = "P num blocks less than D"
1511
continue
1512
1342
- if not local_block_ids:
1513
+ if not any(local_block_ids_by_group):
1514
continue
1515
1345
- # Group by indices
1346
- group_local_block_ids, group_remote_block_ids = group_concurrent_contiguous(
1347
- local_block_ids, remote_block_ids
1516
+ local_block_ids_by_group = self._logical_to_kernel_block_ids(
1517
+ local_block_ids_by_group
1518
+ )
1519
+ remote_block_ids_by_group = self._logical_to_kernel_block_ids(
1520
+ remote_block_ids_by_group
1521
)
1522
1523
for local_region, remote_region in zip(local_regions, remote_regions):
1351
- should_transfer, src_region_offset, dst_region_offset, transfer_len = (
1352
- self._get_sender_transfer_plan(
1353
- local_kv_block_len=local_region.kv_block_len,
1354
- remote_kv_block_len=remote_region.kv_block_len,
1355
- remote_tp_rank=agent_meta.remote_tp_rank,
1356
- remote_tp_size=agent_meta.remote_tp_size,
1357
- )
1524
+ assert local_region.group_index == remote_region.group_index, (
1525
+ "Aligned Mooncake transfer regions must belong to the same "
1526
+ "KV group."
1527
+ )
1528
+ group_index = local_region.group_index
1529
+ assert group_index < len(local_block_ids_by_group), (
1530
+ "Transfer region references a missing KV group."
1531
+ )
1532
+ local_block_ids = local_block_ids_by_group[group_index]
1533
+ remote_block_ids = remote_block_ids_by_group[group_index]
1534
+ if not local_block_ids:
1535
+ continue
1536
+
1537
+ # Group by indices within this region's KV-cache group only.
1538
+ group_local_block_ids, group_remote_block_ids = (
1539
+ group_concurrent_contiguous(local_block_ids, remote_block_ids)
1540
+ )
1541
+ (
1542
+ should_transfer,
1543
+ src_region_offset,
1544
+ dst_region_offset,
1545
+ transfer_len,
1546
+ ) = self._get_sender_transfer_plan(
1547
+ local_kv_block_len=local_region.kv_block_len,
1548
+ remote_kv_block_len=remote_region.kv_block_len,
1549
+ remote_tp_rank=agent_meta.remote_tp_rank,
1550
+ remote_tp_size=agent_meta.remote_tp_size,
1551
)
1552
if not should_transfer:
1553
# Replicated KV cache: only one producer rank in the TP group
@@ -1368,7 +1561,7 @@ class MooncakeConnectorWorker:
1561
"Computed source transfer region exceeds local KV block size."
1562
)
1563
assert dst_region_offset + transfer_len <= remote_region.kv_block_len, (
1371
- "Computed destination transfer region exceeds remote KV block size."
1564
+ "Destination transfer region exceeds remote KV block size."
1565
)
1566
# Collapse one contiguous block group into a single larger
1567
# transfer descriptor when the per-block copy is identical.
@@ -1411,28 +1604,10 @@ class MooncakeConnectorWorker:
1604
)
1605
lengths.append(transfer_len)
1606
1414
- if local_region is local_regions[0]:
1415
- logger.debug(
1416
- "Mooncake transfer plan for request %s: local_tp=%d "
1417
- "remote_tp=%d remote_tp_rank=%d local_block_len=%d "
1418
- "remote_block_len=%d src_offset=%d dst_offset=%d "
1419
- "transfer_len=%d coalesce=%s",
1420
- d_req_id,
1421
- self.tp_size,
1422
- agent_meta.remote_tp_size,
1423
- agent_meta.remote_tp_rank,
1424
- local_region.block_len,
1425
- remote_region.block_len,
1426
- src_region_offset,
1427
- dst_region_offset,
1428
- transfer_len,
1429
- can_coalesce,
1430
- )
1431
-
1607
logger.debug(
1608
"Sending kv_caches for request %s (%d blocks) to %s",
1609
d_req_id,
1435
- len(local_block_ids),
1610
+ sum(len(group) for group in local_block_ids_by_group),
1611
remote_session,
1612
)
1613
@@ -1480,18 +1655,33 @@ class MooncakeConnectorWorker:
1655
1656
logger.info("Registering KV_Caches. use_mla: %s", self.use_mla)
1657
1483
- kv_data_ptrs = []
1484
- kv_data_lens = []
1485
- seen_base_addresses = []
1658
+ kv_data_ptrs: list[int] = []
1659
+ kv_data_lens: list[int] = []
1660
+ region_base_addresses: list[int] = []
1661
+ seen_storage_ptrs: set[int] = set()
1662
self.block_len_per_layer = []
1663
+ self.kv_block_len_per_layer = []
1664
self.registered_layer_names = []
1665
self.registered_layer_indices = []
1666
+ self.registered_group_indices = []
1667
1490
- split_k_and_v = self.transfer_topo.split_k_and_v
1491
- tensor_size_bytes = None
1668
for layer_name, cache_or_caches in kv_caches.items():
1669
layer_index = extract_layer_index(layer_name)
1494
- cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1670
+ layer_spec = self._layer_specs.get(layer_name)
1671
+ if layer_spec is None:
1672
+ logger.debug(
1673
+ "Skipping layer %s because no KV cache spec is present.",
1674
+ layer_name,
1675
+ )
1676
+ continue
1677
+ if isinstance(layer_spec, MambaSpec):
1678
+ conv, _ = cache_or_caches
1679
+ cache_list = [conv]
1680
+ else:
1681
+ cache_list = self.transfer_topo.get_transfer_cache_regions(
1682
+ cache_or_caches, layer_spec
1683
+ )
1684
+
1685
logger.debug(
1686
"registering layer %s with %d cache tensor(s)",
1687
layer_name,
@@ -1501,45 +1691,46 @@ class MooncakeConnectorWorker:
1691
for cache in cache_list:
1692
self._log_debug_cache_registration(layer_name, cache)
1693
base_addr = cache.data_ptr()
1504
- if base_addr in seen_base_addresses:
1505
- continue
1506
-
1507
- seen_base_addresses.append(base_addr)
1508
-
1509
- if tensor_size_bytes is None:
1510
- tensor_size_bytes = cache.nbytes
1511
- self.num_blocks = cache.shape[0]
1512
- assert cache.shape[0] == self.num_blocks, (
1513
- "All kv cache tensors must have the same number of blocks"
1514
- )
1515
-
1516
- # Use stride-based block length so RDMA reaches the last
1517
- # block's padding (e.g. DeepseekV4 MLA alignment). stride(0)
1518
- # reflects the actual byte distance between consecutive
1519
- # blocks in GPU memory, which matches or exceeds the
1520
- # shape-based size.
1694
block_len = cache.stride(0) * cache.element_size()
1695
+ region_base_addresses.append(base_addr)
1696
1697
+ if isinstance(layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)):
1698
+ kv_block_len = layer_spec.page_size_bytes
1699
+ elif self.transfer_topo.virtually_split_kv_in_blocks and not isinstance(
1700
+ layer_spec, MambaSpec
1701
+ ):
1702
+ kv_block_len = block_len // 2
1703
+ else:
1704
+ kv_block_len = block_len
1705
self.block_len_per_layer.append(block_len)
1706
+ self.kv_block_len_per_layer.append(kv_block_len)
1707
self.registered_layer_names.append(layer_name)
1708
self.registered_layer_indices.append(layer_index)
1526
- kv_data_ptrs.append(base_addr)
1527
- kv_data_lens.append(self.num_blocks * block_len)
1709
+ self.registered_group_indices.append(
1710
+ self._layer_group_indices[layer_name]
1711
+ )
1712
+ storage = cache.untyped_storage()
1713
+ storage_addr = storage.data_ptr()
1714
+ if storage_addr not in seen_storage_ptrs:
1715
+ seen_storage_ptrs.add(storage_addr)
1716
+ kv_data_ptrs.append(storage_addr)
1717
+ kv_data_lens.append(storage.nbytes())
1718
1529
- self.kv_caches_base_addr = seen_base_addresses
1530
- self.seen_base_addresses = seen_base_addresses
1719
+ self.kv_caches_base_addr = region_base_addresses
1720
+ self.seen_base_addresses = kv_data_ptrs
1721
+
1722
+ if not kv_data_ptrs:
1723
+ raise RuntimeError("No KV cache tensors were registered with Mooncake.")
1724
1725
ret_value = self.engine.batch_register_memory(kv_data_ptrs, kv_data_lens)
1726
if ret_value != 0:
1727
raise RuntimeError("Mooncake batch memory registration failed.")
1728
1536
- assert tensor_size_bytes is not None
1537
- assert self.num_blocks != 0
1729
self.device_kv_caches = kv_caches
1730
logger.debug(
1540
- "registered num_blocks=%d block_lens=%s",
1541
- self.num_blocks,
1731
+ "registered block_lens=%s kv_block_lens=%s",
1732
self.block_len_per_layer,
1733
+ self.kv_block_len_per_layer,
1734
)
1735
1736
# No need to launch server for D node.
@@ -1642,8 +1833,10 @@ class MooncakeConnectorWorker:
1833
},
1834
kv_caches_base_addr=self.kv_caches_base_addr,
1835
block_lens=self.block_len_per_layer,
1836
+ kv_block_lens=self.kv_block_len_per_layer,
1837
registered_layer_names=self.registered_layer_names,
1838
registered_layer_indices=self.registered_layer_indices,
1839
+ registered_group_indices=self.registered_group_indices,
1840
)
1841
1842
encoded_data = self._encoder.encode(metadata)
@@ -1852,15 +2045,34 @@ class MooncakeConnectorWorker:
2045
self,
2046
base_addrs: list[int],
2047
block_lens: list[int],
2048
+ kv_block_lens: list[int],
2049
layer_names: list[str],
2050
layer_indices: list[int],
2051
+ group_indices: list[int] | None = None,
2052
) -> list[TransferRegion]:
2053
+ if not group_indices:
2054
+ group_indices = [
2055
+ self._layer_group_indices.get(layer_name, 0)
2056
+ for layer_name in layer_names
2057
+ ]
2058
+ split_kv_regions = None
2059
+ if self.transfer_topo.virtually_split_kv_in_blocks:
2060
+ split_kv_regions = [
2061
+ not isinstance(
2062
+ self._layer_specs[layer_name],
2063
+ (MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec),
2064
+ )
2065
+ for layer_name in layer_names
2066
+ ]
2067
return _expand_transfer_regions(
2068
base_addrs=base_addrs,
2069
block_lens=block_lens,
2070
+ kv_block_lens=kv_block_lens,
2071
layer_names=layer_names,
2072
layer_indices=layer_indices,
2073
is_kv_layout_blocks_first=self.transfer_topo.virtually_split_kv_in_blocks,
2074
+ group_indices=group_indices,
2075
+ split_kv_regions=split_kv_regions,
2076
)
2077
2078
def _get_sender_transfer_plan(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+1
-1
@@ -320,7 +320,7 @@ class NixlBaseConnectorScheduler:
320
logger.warning("Connection listener got unexpected message %s", msg)
321
sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))
322
323
- def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int:
323
+ def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
324
"""D-side only. Returns N-1 for Mamba models since the decoder
325
always recomputes the last token and must start from h(N-1)."""
326
if self._has_mamba and num_prompt_tokens > 1:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+1
-1
@@ -60,7 +60,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
60
if params is not None and params.get("do_remote_prefill"):
61
# Remote prefill: get all prompt blocks from remote.
62
token_ids = request.prompt_token_ids or []
63
- actual = self._mamba_prefill_token_count(len(token_ids))
63
+ actual = self._get_remote_prefill_token_count(len(token_ids))
64
count = actual - num_computed_tokens
65
if count > 0:
66
return count, True
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
+1
-1
@@ -116,7 +116,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
116
117
if params is not None and params.get("do_remote_prefill"):
118
token_ids = request.prompt_token_ids or []
119
- actual = self._mamba_prefill_token_count(len(token_ids))
119
+ actual = self._get_remote_prefill_token_count(len(token_ids))
120
count = actual - num_computed_tokens
121
if count > 0:
122
return count, True