[Perf][KVConnector][Mooncake] Vectorize prepare_value on the KV load path (#48531)
Signed-off-by: girasoley <girasoley@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: girasoley <girasoley@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Summer Yang committed
Jul 23, 2026 at 03:12 UTC
a4904ba9032af87610965ba7c90267bf7b21b425
4 files changed
+211
-21
tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for ChunkedTokenDatabase.prepare_values."""
4
+
5
+import random
6
+
7
+import pytest
8
+
9
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
10
+ ChunkedTokenDatabase,
11
+ KeyMetadata,
12
+)
13
+from vllm.utils.math_utils import cdiv
14
+
15
+BLOCK_SIZE = 128
16
+
17
+
18
+def _reference_prepare_value(
19
+ db: ChunkedTokenDatabase, start: int, end: int, block_ids: list[int]
20
+) -> tuple[list[int], list[int], int]:
21
+ """Compute a token range with the original scalar implementation."""
22
+ addr_list = []
23
+ size_list = []
24
+ block_id = block_ids[start // db.block_size]
25
+ length = len(db.block_len)
26
+ for index, base_addr in enumerate(db.kv_caches_base_addr):
27
+ addr = base_addr + block_id * db.block_len[index % length]
28
+ assert (end - start) % db.block_size == 0
29
+ size = db.block_len[index % length] * cdiv(end - start, db.block_size)
30
+ addr_list.append(addr)
31
+ size_list.append(size)
32
+ return addr_list, size_list, block_id
33
+
34
+
35
+def _make_db(num_regions: int, num_block_lens: int) -> ChunkedTokenDatabase:
36
+ md = KeyMetadata(model_name="t", tp_rank=1, pcp_rank=0, dcp_rank=0, pp_rank=0)
37
+ db = ChunkedTokenDatabase(md, BLOCK_SIZE)
38
+ db.set_kv_caches_base_addr(
39
+ [0x7F00_0000_0000 + i * (1 << 30) for i in range(num_regions)]
40
+ )
41
+ # Exercise repeated block lengths when there are more cache regions.
42
+ db.set_block_len([30_208 + 512 * i for i in range(num_block_lens)])
43
+ return db
44
+
45
+
46
+@pytest.mark.parametrize("num_regions,num_block_lens", [(96, 96), (96, 2), (1, 1)])
47
+def test_prepare_values_matches_reference(num_regions: int, num_block_lens: int):
48
+ db = _make_db(num_regions, num_block_lens)
49
+ rng = random.Random(0)
50
+ n_blocks = 300
51
+ block_ids = [rng.randrange(0, 1 << 20) for _ in range(n_blocks)]
52
+ chunks = []
53
+ b = 0
54
+ while b < n_blocks - 4:
55
+ span = rng.choice([1, 1, 1, 2, 4])
56
+ chunks.append((b * BLOCK_SIZE, (b + span) * BLOCK_SIZE))
57
+ b += span + rng.choice([0, 1])
58
+
59
+ addrs, sizes, bids = db.prepare_values(chunks, block_ids)
60
+ assert len(addrs) == len(sizes) == len(bids) == len(chunks)
61
+ for (start, end), addr, size, bid in zip(chunks, addrs, sizes, bids):
62
+ ref_addr, ref_size, ref_bid = _reference_prepare_value(
63
+ db, start, end, block_ids
64
+ )
65
+ assert addr == ref_addr
66
+ assert size == ref_size
67
+ assert bid == ref_bid
68
+ # Native bindings require Python ints rather than numpy scalars.
69
+ assert all(type(a) is int for a in addr)
70
+ assert type(bid) is int
71
+
72
+
73
+def test_prepare_value_single_matches_reference():
74
+ db = _make_db(8, 8)
75
+ block_ids = list(range(64))
76
+ got = db.prepare_value(5 * BLOCK_SIZE, 7 * BLOCK_SIZE, block_ids)
77
+ assert got == _reference_prepare_value(
78
+ db, 5 * BLOCK_SIZE, 7 * BLOCK_SIZE, block_ids
79
+ )
80
+
81
+
82
+def test_prepare_values_empty():
83
+ db = _make_db(4, 4)
84
+ assert db.prepare_values([], [1, 2, 3]) == ([], [], [])
85
+
86
+
87
+def test_prepare_values_rejects_unaligned_chunk():
88
+ db = _make_db(4, 4)
89
+ with pytest.raises(AssertionError):
90
+ db.prepare_values([(0, BLOCK_SIZE + 1)], [0, 1])
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+62
@@ -655,6 +655,68 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks():
655
assert masked_hashes == [b"a2".hex()]
656
657
658
+def test_store_sending_thread_prepares_missing_chunks_once_per_group():
659
+ store = MagicMock()
660
+ store.batch_is_exist.return_value = [0, 1, 0, 1, 0, 0]
661
+ store.batch_put_from_multi_buffers.return_value = [256, 256, 512, 512]
662
+ coord = SimpleNamespace(
663
+ lcm_block_size=16,
664
+ store_mask=lambda token_len, start_token, num_prompt_tokens=None: (
665
+ None,
666
+ None,
667
+ ),
668
+ )
669
+
670
+ db0 = ChunkedTokenDatabase(
671
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
672
+ block_size=16,
673
+ )
674
+ db0.set_kv_caches_base_addr([0x1000])
675
+ db0.set_block_len([256])
676
+ db0.prepare_values = MagicMock(wraps=db0.prepare_values)
677
+ db0.prepare_value = MagicMock(side_effect=AssertionError("scalar path called"))
678
+
679
+ db1 = ChunkedTokenDatabase(
680
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
681
+ block_size=16,
682
+ )
683
+ db1.set_kv_caches_base_addr([0x2000])
684
+ db1.set_block_len([512])
685
+ db1.prepare_values = MagicMock(wraps=db1.prepare_values)
686
+ db1.prepare_value = MagicMock(side_effect=AssertionError("scalar path called"))
687
+
688
+ thread = _make_store_sending_thread(
689
+ store,
690
+ coord=coord,
691
+ token_databases=[db0, db1],
692
+ )
693
+ thread.add_stored_request("req-a")
694
+ thread._handle_request(
695
+ ReqMeta(
696
+ req_id="req-a",
697
+ token_len_chunk=48,
698
+ block_ids=([0, 1, 2], [2, 1, 0]),
699
+ block_hashes=[b"a0", b"a1", b"a2"],
700
+ can_save=True,
701
+ )
702
+ )
703
+
704
+ db0.prepare_value.assert_not_called()
705
+ db1.prepare_value.assert_not_called()
706
+ db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [0, 1, 2])
707
+ db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [2, 1, 0])
708
+
709
+ keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args
710
+ assert [key.rsplit("@", 1)[-1] for key in keys] == [
711
+ "6130",
712
+ "6132",
713
+ "6131",
714
+ "6132",
715
+ ]
716
+ assert addrs == [[0x1000], [0x1200], [0x2200], [0x2000]]
717
+ assert sizes == [[256], [256], [512], [512]]
718
+
719
+
720
def test_store_sending_thread_only_skips_on_no_available_handle():
721
store = MagicMock()
722
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+36
-11
@@ -9,6 +9,7 @@ from collections.abc import Iterable, Sequence
9
from dataclasses import dataclass
10
from typing import cast
11
12
+import numpy as np
13
import torch
14
15
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -196,22 +197,46 @@ class ChunkedTokenDatabase:
197
def prepare_value(
198
self, start: int, end: int, block_ids: list[int]
199
) -> tuple[list[int], list[int], int]:
199
- """Compute memory addresses and sizes for a token range.
200
+ """Compute memory addresses and sizes for a single token range.
201
202
Returns:
203
(addr_list, size_list, block_id)
204
"""
204
- addr_list = []
205
- size_list = []
206
- block_id = block_ids[start // self.block_size]
205
+ addr_lists, size_lists, chunk_block_ids = self.prepare_values(
206
+ ((start, end),), block_ids
207
+ )
208
+ return addr_lists[0], size_lists[0], chunk_block_ids[0]
209
+
210
+ def prepare_values(
211
+ self,
212
+ chunks: Sequence[tuple[int, int]],
213
+ block_ids: list[int],
214
+ ) -> tuple[list[list[int]], list[list[int]], list[int]]:
215
+ """Compute memory addresses and sizes for multiple token ranges.
216
+
217
+ Returns:
218
+ (addr_lists, size_lists, chunk_block_ids), one entry per chunk.
219
+ """
220
+ if not chunks:
221
+ return [], [], []
222
+ base = np.asarray(self.kv_caches_base_addr, dtype=np.int64)
223
length = len(self.block_len)
208
- for index, base_addr in enumerate(self.kv_caches_base_addr):
209
- addr = base_addr + block_id * self.block_len[index % length]
210
- assert (end - start) % self.block_size == 0
211
- size = self.block_len[index % length] * cdiv(end - start, self.block_size)
212
- addr_list.append(addr)
213
- size_list.append(size)
214
- return addr_list, size_list, block_id
224
+ blen = np.asarray(
225
+ [self.block_len[i % length] for i in range(base.shape[0])],
226
+ dtype=np.int64,
227
+ )
228
+ n = len(chunks)
229
+ starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n)
230
+ spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts
231
+ assert not (spans % self.block_size).any()
232
+ bids = np.fromiter(
233
+ (block_ids[i] for i in (starts // self.block_size).tolist()),
234
+ dtype=np.int64,
235
+ count=n,
236
+ )
237
+ addrs = base[None, :] + bids[:, None] * blen[None, :]
238
+ sizes = blen[None, :] * (spans // self.block_size)[:, None]
239
+ return addrs.tolist(), sizes.tolist(), bids.tolist()
240
241
def process_tokens(
242
self,
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+23
-10
@@ -634,6 +634,21 @@ class KVCacheStoreSendingThread(KVTransferThread):
634
addrs: list[list[int]] = []
635
sizes: list[list[int]] = []
636
stored_events: list[BlockStored] = []
637
+ chunks_per_group: list[list[tuple[int, int]]] = [
638
+ [] for _ in self.token_databases
639
+ ]
640
+ for start, end, g_idx in zip(starts, ends, group_indices, strict=True):
641
+ chunks_per_group[g_idx].append((start, end))
642
+ for g_idx, chunks in enumerate(chunks_per_group):
643
+ if not chunks:
644
+ continue
645
+ db = self.token_databases[g_idx]
646
+ group_addrs, group_sizes, _ = db.prepare_values(
647
+ chunks, block_ids_per_group[g_idx]
648
+ )
649
+ addrs.extend(group_addrs)
650
+ sizes.extend(group_sizes)
651
+
652
# parent_block_hash chains live within a group, not across.
653
if self.enable_kv_event:
654
prev_key_per_group: dict[int, Any] = {}
@@ -645,10 +660,6 @@ class KVCacheStoreSendingThread(KVTransferThread):
660
zip(starts, ends, group_indices, strict=True)
661
):
662
db = self.token_databases[g_idx]
648
- addr, size, _ = db.prepare_value(s, e, block_ids_per_group[g_idx])
649
- addrs.append(addr)
650
- sizes.append(size)
651
-
663
if self.enable_kv_event:
664
token_ids = (
665
req_meta.token_ids[s:e]
@@ -805,19 +816,21 @@ class KVCacheStoreRecvingThread(KVTransferThread):
816
block_id_list: list[int] = []
817
for g_idx, db in enumerate(self.token_databases):
818
mask = load_mask_per_group[g_idx]
819
+ chunks: list[tuple[int, int]] = []
820
for start, end, block_hash in db.process_tokens(
821
token_len, req_meta.block_hashes, mask_num
822
):
823
chunk_idx = start // db.block_size
824
if chunk_idx >= len(mask) or not mask[chunk_idx]:
825
continue
814
- addr, size, block_id = db.prepare_value(
815
- start, end, req_meta.block_ids[g_idx]
816
- )
826
key_list.append(db.key_for(block_hash))
818
- addr_list.append(addr)
819
- size_list.append(size)
820
- block_id_list.append(block_id)
827
+ chunks.append((start, end))
828
+ g_addrs, g_sizes, g_block_ids = db.prepare_values(
829
+ chunks, req_meta.block_ids[g_idx]
830
+ )
831
+ addr_list.extend(g_addrs)
832
+ size_list.extend(g_sizes)
833
+ block_id_list.extend(g_block_ids)
834
835
# Rotate aligned lists by tp_rank for load balancing.
836
rotation = self.tp_rank % len(key_list)