[Bugfix][KV Connector] Fix Mooncake physical-block transfer length (#54272)
Signed-off-by: zhewenl <zhewenl@users.noreply.github.com> Co-authored-by: zhewenl <zhewenl@users.noreply.github.com>
Zhewen Li committed
Sep 1, 2026 at 13:23 UTC
5cc32fbdfac44987bc2ece1d9ded7cb43acc2820
3 files changed
+208
-7
tests/v1/kv_connector/unit/test_mooncake_connector.py
+3
-2
@@ -1215,8 +1215,9 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
1215
worker.use_mla = True
1216
worker.transfer_topo.is_mla = True
1217
1218
- # MLA cache tensor: shape[-2] is the block size.
1219
- mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16)
1218
+ # MLA cache tensor: shape[-2] is the block size and each block's
1219
+ # byte stride matches the cache spec page size.
1220
+ mla_cache = torch.zeros((2, 16, 512), dtype=torch.float16)
1221
# Eagle3/GQA-like cache tensor: shape[-2] is num_kv_heads, not block size.
1222
eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16)
1223
kv_caches = {
tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
+187
@@ -199,6 +199,193 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
199
connector.connector_worker = None
200
201
202
+def test_register_kv_caches_scales_attention_len_to_kernel_block(monkeypatch):
203
+ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
204
+ vllm_config = create_vllm_config(
205
+ kv_connector="MooncakeConnector",
206
+ kv_role="kv_consumer",
207
+ )
208
+ kv_cache_config = make_hybrid_gdn_kv_cache_config(
209
+ vllm_config.cache_config.block_size
210
+ )
211
+
212
+ with set_current_vllm_config(vllm_config), patch_worker_dependencies():
213
+ connector = MooncakeConnector(
214
+ vllm_config,
215
+ KVConnectorRole.WORKER,
216
+ kv_cache_config,
217
+ )
218
+ worker = connector.connector_worker
219
+ factor = 4
220
+ worker._physical_blocks_per_logical_kv_block = factor
221
+
222
+ fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
223
+ physical_page_bytes = fa_spec.page_size_bytes // factor
224
+ fa_cache = torch.empty(
225
+ kv_cache_config.num_blocks * factor,
226
+ physical_page_bytes,
227
+ dtype=torch.uint8,
228
+ )
229
+
230
+ worker.register_kv_caches({"model.layers.0.self_attn": fa_cache})
231
+
232
+ assert worker.block_len_per_layer == [physical_page_bytes]
233
+ assert worker.kv_block_len_per_layer == [physical_page_bytes]
234
+
235
+ worker.shutdown()
236
+ worker.shutdown = noop_shutdown
237
+ connector.connector_worker = None
238
+
239
+
240
+@pytest.mark.asyncio
241
+async def test_build_transfer_params_uses_non_overlapping_physical_pages():
242
+ kv_cache_config = make_hybrid_gdn_kv_cache_config(block_size=16)
243
+ worker = object.__new__(MooncakeConnectorWorker)
244
+ worker.shutdown = noop_shutdown
245
+ worker.use_mla = False
246
+ worker.engine = SimpleNamespace(batch_register_memory=lambda *_: 0)
247
+ worker.is_kv_consumer = True
248
+ worker.tp_rank = 0
249
+ worker.tp_size = 1
250
+ worker.transfer_topo = SimpleNamespace(local_replicates_kv_cache=False)
251
+ worker.kv_cache_config = kv_cache_config
252
+ worker._layer_specs = {
253
+ layer_name: group.kv_cache_spec
254
+ for group in kv_cache_config.kv_cache_groups
255
+ for layer_name in group.layer_names
256
+ }
257
+ worker._layer_group_indices = {
258
+ layer_name: group_index
259
+ for group_index, group in enumerate(kv_cache_config.kv_cache_groups)
260
+ for layer_name in group.layer_names
261
+ }
262
+
263
+ ratio = 4
264
+ worker._physical_blocks_per_logical_kv_block = ratio
265
+ fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
266
+ physical_page_bytes = fa_spec.page_size_bytes // ratio
267
+ fa_cache = torch.empty(
268
+ kv_cache_config.num_blocks * ratio,
269
+ physical_page_bytes,
270
+ dtype=torch.uint8,
271
+ )
272
+ worker.register_kv_caches({"model.layers.0.self_attn": fa_cache})
273
+
274
+ local_regions = worker._get_transfer_regions(
275
+ worker.kv_caches_base_addr,
276
+ worker.block_len_per_layer,
277
+ worker.kv_block_len_per_layer,
278
+ worker.registered_layer_names,
279
+ worker.registered_layer_indices,
280
+ worker.registered_group_indices,
281
+ )
282
+ remote_base_addr = 0x100000
283
+ remote_regions = [
284
+ TransferRegion(
285
+ layer_name="model.layers.0.self_attn",
286
+ layer_index=0,
287
+ base_addr=remote_base_addr,
288
+ block_len=physical_page_bytes,
289
+ kv_block_len=physical_page_bytes,
290
+ group_index=0,
291
+ )
292
+ ]
293
+
294
+ local_logical_block = 2
295
+ remote_logical_block = 5
296
+ transfer_id = "xfer-physical-attention-pages"
297
+ send_meta = SendBlockMeta(
298
+ p_req_id="p-physical-attention-pages",
299
+ transfer_id=transfer_id,
300
+ local_block_ids=[[local_logical_block], []],
301
+ ready=asyncio.Event(),
302
+ )
303
+ xfer_meta = MooncakeXferMetadata(
304
+ remote_hostname="consumer-host",
305
+ remote_port=54321,
306
+ remote_tp_size=1,
307
+ remote_tp_rank=0,
308
+ req_blocks={
309
+ "d-physical-attention-pages": (
310
+ transfer_id,
311
+ [[remote_logical_block], []],
312
+ )
313
+ },
314
+ kv_caches_base_addr=[],
315
+ block_lens=[],
316
+ kv_block_lens=[],
317
+ )
318
+
319
+ (
320
+ src_ptrs,
321
+ dst_ptrs,
322
+ lengths,
323
+ err_reqs,
324
+ err_msg,
325
+ ) = await worker._build_transfer_params(
326
+ [("d-physical-attention-pages", send_meta)],
327
+ xfer_meta,
328
+ local_regions,
329
+ remote_regions,
330
+ )
331
+
332
+ assert err_reqs == []
333
+ assert err_msg is None
334
+ assert src_ptrs == [
335
+ fa_cache.data_ptr() + local_logical_block * ratio * physical_page_bytes
336
+ ]
337
+ assert dst_ptrs == [
338
+ remote_base_addr + remote_logical_block * ratio * physical_page_bytes
339
+ ]
340
+ assert lengths == [ratio * physical_page_bytes]
341
+
342
+ def split_contiguous_blocks(src_ids, dst_ids):
343
+ return [[block_id] for block_id in src_ids], [
344
+ [block_id] for block_id in dst_ids
345
+ ]
346
+
347
+ with patch(
348
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
349
+ "mooncake_connector.group_concurrent_contiguous",
350
+ side_effect=split_contiguous_blocks,
351
+ ):
352
+ (
353
+ src_ptrs,
354
+ dst_ptrs,
355
+ lengths,
356
+ err_reqs,
357
+ err_msg,
358
+ ) = await worker._build_transfer_params(
359
+ [("d-physical-attention-pages", send_meta)],
360
+ xfer_meta,
361
+ local_regions,
362
+ remote_regions,
363
+ )
364
+
365
+ expected_src_ptrs = [
366
+ fa_cache.data_ptr()
367
+ + (local_logical_block * ratio + offset) * physical_page_bytes
368
+ for offset in range(ratio)
369
+ ]
370
+ expected_dst_ptrs = [
371
+ remote_base_addr + (remote_logical_block * ratio + offset) * physical_page_bytes
372
+ for offset in range(ratio)
373
+ ]
374
+ assert err_reqs == []
375
+ assert err_msg is None
376
+ assert src_ptrs == expected_src_ptrs
377
+ assert dst_ptrs == expected_dst_ptrs
378
+ assert lengths == [physical_page_bytes] * ratio
379
+ assert all(
380
+ next_ptr - ptr == physical_page_bytes
381
+ for ptr, next_ptr in zip(src_ptrs, src_ptrs[1:])
382
+ )
383
+ assert all(
384
+ next_ptr - ptr == physical_page_bytes
385
+ for ptr, next_ptr in zip(dst_ptrs, dst_ptrs[1:])
386
+ )
387
+
388
+
389
def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
390
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
391
vllm_config = create_vllm_config(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+18
-5
@@ -1675,11 +1675,24 @@ class MooncakeConnectorWorker:
1675
block_len = region_cache.stride(0) * region_cache.element_size()
1676
region_base_addresses.append(base_addr)
1677
1678
- kv_block_len = (
1679
- layer_spec.page_size_bytes
1680
- if isinstance(layer_spec, AttentionSpec) and block_is_contiguous
1681
- else block_len
1682
- )
1678
+ if isinstance(layer_spec, AttentionSpec) and block_is_contiguous:
1679
+ assert (
1680
+ layer_spec.page_size_bytes
1681
+ % self._physical_blocks_per_logical_kv_block
1682
+ == 0
1683
+ )
1684
+ kv_block_len = (
1685
+ layer_spec.page_size_bytes
1686
+ // self._physical_blocks_per_logical_kv_block
1687
+ )
1688
+ else:
1689
+ kv_block_len = block_len
1690
+ if kv_block_len > block_len:
1691
+ raise RuntimeError(
1692
+ "Mooncake transfer length exceeds physical block stride "
1693
+ f"for {layer_name}: kv_block_len={kv_block_len}, "
1694
+ f"block_len={block_len}."
1695
+ )
1696
self.block_len_per_layer.append(block_len)
1697
self.kv_block_len_per_layer.append(kv_block_len)
1698
self.registered_layer_names.append(layer_name)