@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c866ba9d11

[KV Connector] Support heterogeneous TP sharing in Mooncake Store Connector (#53129)

Signed-off-by: Guanyi Chen <939416532@qq.com> Signed-off-by: z-zanez <zhouzh93@mail2.sysu.edu.cn> Signed-off-by: Zheng Zhou <zhouzh93@mail2.sysu.edu.cn> Signed-off-by: zane <zhouzh93@mail2.sysu.edu.cn> Co-authored-by: Guanyi Chen <939416532@qq.com>

周峥 committed Sep 2, 2026 at 01:18 UTC c866ba9d11984fd621af5837f91cd763312f4dd3

8 files changed +1766 -361

docs/features/mooncake_store_connector_usage.md

+55 -9

@@ -137,15 +137,12 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \

137 ```

138

139 To also offload newly completed decode KV blocks, add the following extra

140 -configuration to the decoder's `MooncakeStoreConnector` entry. This currently

141 -requires the prefill and decode instances to use the same tensor-parallel size

142 -and compatible KV-cache topology.

140 +configuration to the decoder's `MooncakeStoreConnector` entry.

141

142 When decode processing starts, the consumer checks the block-aligned prompt

143 prefix and fills any blocks missing from the Store. Subsequent saves append

144 newly completed decode blocks. This keeps a complete, reusable prefix in the

147 -Store and also covers deployments where prompt KV is delivered directly by

148 -`MooncakeConnector` instead of through the Store.

145 +Store. It also covers prompt KV delivered directly by `MooncakeConnector`.

146

147 ```json

148 {

@@ -155,9 +152,41 @@ Store and also covers deployments where prompt KV is delivered directly by

152 }

153 ```

154

155 +#### Sharing one Store across multiple Prefill TP sizes

156 +

157 +Heterogeneous-TP sharing normally uses a fixed `store_tp_size`. When several

158 +prefillers use different TP sizes, opt in to a common Store TP derived from

159 +their least common multiple:

160 +

161 +```json

162 +{

163 + "kv_connector_extra_config": {

164 + "enable_store_tp_lcm": true,

165 + "prefill_tp_sizes": [4, 2]

166 + }

167 +}

168 +```

169 +

170 +Every prefiller and decoder that shares these entries must use the same list.

171 +The example selects Store TP 4: a TP4 endpoint maps each rank to one Store

172 +shard, while TP2 endpoints map each rank to two Store shards. Runtime TP sizes

173 +remain unchanged. A decoder configured with `"save_decode_cache": true` uses

174 +the same Store TP for decode KV from every prefiller.

175 +

176 +The list may contain positive integer TP sizes. Sharing requires a Store TP that

177 +is at least the local TP and divisible by it, an LBHNC or LBNHC local KV cache,

178 +and the existing topology and KV-head constraints. The Store namespace includes

179 +the attention backend's selected layout. Different layouts use separate Store

180 +entries. Malformed lists and unsupported endpoints use an isolated rank-local

181 +key layout. When

182 +`enable_store_tp_lcm` is absent or false, `prefill_tp_sizes` has no effect and

183 +the existing `store_tp_size` behavior is unchanged.

184 +

185 **Proxy:**

186

160 -A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details.

187 +A disaggregation proxy routes requests between prefiller and decoder nodes.

188 +When `MooncakeConnector` is also used for direct P2P transfer, refer to its

189 +[usage guide](mooncake_connector_usage.md) for proxy setup details.

190

191 ### Disk Offloading

192

@@ -249,9 +278,26 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants

278 - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.

279 - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).

280 - `save_decode_cache` (bool): Enable offloading decode tokens' KV cache. A `kv_consumer` does not save during prefill; when decode starts, it fills any missing block-aligned prompt prefix before appending completed decode blocks. Default: `false`.

252 -

253 -Decode offloading uses the existing TP-rank key namespace. Cross-TP sharing is

254 -not yet supported.

281 +- `store_tp_size` (int): Common Store TP for endpoints with different local TP sizes. It supports LBHNC and LBNHC local KV caches, with `store_tp_size >= local_tp_size` and `store_tp_size % local_tp_size == 0`. The current topology is one full-attention cache group, PCP/DCP disabled, and cross-layer blocks disabled. For GQA and MHA, the total KV-head count must be divisible by `store_tp_size`. Store shards contain fixed global KV-head ranges in the local layout. Shared endpoints use the same KV cache layout, pipeline-parallel size, and Store TP. The Store namespace includes the layout and PP size. Unsupported configurations use a topology-specific rank-local namespace.

282 +

283 +LBHNC/HND is strongly recommended for TP-sharded Store when supported. LBNHC/NHD

284 +creates many transfer segments and may significantly reduce PUT/GET performance.

285 +

286 +For example, with prefill TP 4, decode TP 2, and eight KV heads, set

287 +`store_tp_size` to 4 on both instances. Each decode rank reads and writes two

288 +of the four Store shards.

289 +

290 +MQA with one total KV head uses a replicated-head layout. For the supported

291 +prefill TP 4 to decode TP 2 case, every rank uses the

292 +same rank-0 key namespace. The four prefill replicas stripe block PUTs so each

293 +object is stored once, while both decode ranks GET every block into their local

294 +KV replica. `store_tp_size` does not appear in MQA keys, so identical MQA

295 +objects written at different store TP sizes share the same pool entry when PP

296 +sizes match.

297 +

298 +Tensor-parallel collectives and low-precision arithmetic are not bitwise

299 +invariant across TP sizes, so heterogeneous-TP reuse does not guarantee the

300 +same greedy output as recomputing the prefix at the decode TP size.

301

302 ## Notes

303

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

+170 -4

@@ -48,6 +48,31 @@ def _make_vllm_config():

48 )

51 +@pytest.mark.parametrize(

52 + "extra_config",

53 + [

54 + {"store_tp_size": 4},

55 + {"enable_store_tp_lcm": True, "prefill_tp_sizes": [6, 4]},

56 + {"prefill_tp_sizes": [4, 2]},

57 + {"enable_store_tp_lcm": True, "prefill_tp_sizes": [4, 0]},

58 + None,

59 + ],

60 +)

61 +def test_store_tp_does_not_override_backend_kv_cache_layout(extra_config):

62 + vllm_config = create_vllm_config(

63 + kv_connector="MooncakeStoreConnector",

64 + kv_role="kv_both",

65 + kv_connector_extra_config=extra_config,

66 + )

67 +

68 + assert (

69 + mooncake_store_connector.MooncakeStoreConnector.get_required_kvcache_layout(

70 + vllm_config

71 + )

72 + is None

73 + )

74 +

75 +

76 def _make_kv_cache_config() -> KVCacheConfig:

77 """Single-group full-attention KVCacheConfig — enough for the connector

78 constructor's validate() pass."""

@@ -90,15 +115,18 @@ def test_scheduler_requires_align_mode_for_mamba():

115 scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)

116

117

93 -def _make_block_stored() -> BlockStored:

118 +def _make_block_stored(

119 + block_hash: bytes = b"hash", group_idx: int | None = None

120 +) -> BlockStored:

121 return BlockStored(

95 - block_hashes=[b"hash"],

122 + block_hashes=[block_hash],

123 parent_block_hash=None,

124 token_ids=[1, 2, 3],

125 block_size=16,

126 lora_id=None,

127 medium="cpu",

128 lora_name=None,

129 + group_idx=group_idx,

130 )

131

132

@@ -169,7 +197,7 @@ def test_worker_methods_delegate_to_store_worker():

197 assert invalid_block_ids == {3, 4}

198

199

172 -def test_get_kv_connector_kv_cache_events_returns_none_when_empty():

200 +def test_get_kv_connector_kv_cache_events_returns_none_when_disabled():

201 vllm_config = _make_vllm_config()

202 kv_cache_config = _make_kv_cache_config()

203

@@ -184,7 +212,7 @@ def test_get_kv_connector_kv_cache_events_returns_none_when_empty():

212 vllm_config, KVConnectorRole.WORKER, kv_cache_config

213 )

214

187 - mock_worker_cls.return_value.get_kv_events.return_value = []

215 + mock_worker_cls.return_value.enable_kv_events = False

216 assert connector.get_kv_connector_kv_cache_events() is None

217

218

@@ -248,6 +276,9 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events():

276 )

277

278 mock_worker_cls.return_value.get_kv_events.return_value = [event]

279 + mock_worker_cls.return_value.enable_kv_events = True

280 + mock_worker_cls.return_value.kv_send_thread = MagicMock()

281 + mock_worker_cls.return_value.group_tp_replication_factors = (1,)

282 kv_events = connector.get_kv_connector_kv_cache_events()

283

284 assert isinstance(kv_events, mooncake_store_connector.MooncakeStoreKVEvents)

@@ -255,6 +286,119 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events():

286 assert kv_events.get_all_events() == [event]

287

288

289 +def test_get_kv_connector_kv_cache_events_keeps_empty_worker_contribution():

290 + vllm_config = _make_vllm_config()

291 + kv_cache_config = _make_kv_cache_config()

292 +

293 + with (

294 + set_current_vllm_config(vllm_config),

295 + patch(

296 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."

297 + "connector.MooncakeStoreWorker"

298 + ) as mock_worker_cls,

299 + ):

300 + connector = mooncake_store_connector.MooncakeStoreConnector(

301 + vllm_config, KVConnectorRole.WORKER, kv_cache_config

302 + )

303 +

304 + worker = mock_worker_cls.return_value

305 + worker.enable_kv_events = True

306 + worker.kv_send_thread = MagicMock()

307 + worker.group_tp_replication_factors = (2,)

308 + worker.get_kv_events.return_value = []

309 +

310 + kv_events = connector.get_kv_connector_kv_cache_events()

311 +

312 + assert isinstance(kv_events, mooncake_store_connector.MooncakeStoreKVEvents)

313 + assert kv_events.get_number_of_workers() == 1

314 + assert kv_events.get_all_events() == []

315 +

316 +

317 +def _aggregate_store_events(

318 + worker_events: list[list[BlockStored]],

319 + replication_factors: tuple[int, ...],

320 +) -> list[BlockStored]:

321 + combined = mooncake_store_connector.MooncakeStoreKVEvents(

322 + num_workers=1,

323 + group_tp_replication_factors=replication_factors,

324 + )

325 + combined.add_events(worker_events[0])

326 + for events in worker_events[1:]:

327 + combined.add_events(events)

328 + combined.increment_workers()

329 + combined.aggregate()

330 + return combined.get_all_events()

331 +

332 +

333 +def test_gqa_store_events_follow_tp_replica_striping():

334 + block_a = _make_block_stored(b"a", group_idx=0)

335 + block_b = _make_block_stored(b"b", group_idx=0)

336 +

337 + events = _aggregate_store_events(

338 + [[block_a], [block_b], [block_a], [block_b]],

339 + replication_factors=(2,),

340 + )

341 +

342 + assert events == [block_a, block_b]

343 +

344 +

345 +def test_mqa_store_events_require_one_replica():

346 + blocks = [_make_block_stored(bytes([index]), group_idx=0) for index in range(4)]

347 +

348 + events = _aggregate_store_events(

349 + [[block] for block in blocks],

350 + replication_factors=(4,),

351 + )

352 +

353 + assert events == blocks

354 +

355 +

356 +def test_store_events_require_every_non_replicated_worker():

357 + complete = _make_block_stored(b"complete", group_idx=0)

358 + incomplete = _make_block_stored(b"incomplete", group_idx=0)

359 +

360 + events = _aggregate_store_events(

361 + [

362 + [complete, incomplete],

363 + [complete],

364 + [complete],

365 + [complete],

366 + ],

367 + replication_factors=(1,),

368 + )

369 +

370 + assert events == [complete]

371 +

372 +

373 +def test_store_events_reject_incomplete_replica_coverage():

374 + complete = _make_block_stored(b"complete", group_idx=0)

375 + incomplete = _make_block_stored(b"incomplete", group_idx=0)

376 +

377 + events = _aggregate_store_events(

378 + [[complete, incomplete], [complete], [], []],

379 + replication_factors=(2,),

380 + )

381 +

382 + assert events == [complete]

383 +

384 +

385 +def test_store_events_use_each_groups_replication_factor():

386 + full_attention = _make_block_stored(b"full", group_idx=0)

387 + replicated_attention = _make_block_stored(b"replicated", group_idx=1)

388 +

389 + events = _aggregate_store_events(

390 + [

391 + [full_attention, replicated_attention],

392 + [full_attention],

393 + [full_attention, replicated_attention],

394 + [full_attention],

395 + ],

396 + replication_factors=(1, 2),

397 + )

398 +

399 + assert events == [full_attention, replicated_attention]

400 +

401 +

402 def test_update_connector_output_and_take_events():

403 vllm_config = _make_vllm_config()

404 kv_cache_config = _make_kv_cache_config()

@@ -284,6 +428,28 @@ def test_update_connector_output_and_take_events():

428 assert connector._kv_cache_events is None

429

430

431 +def test_store_events_persist_across_polls():

432 + connector = object.__new__(mooncake_store_connector.MooncakeStoreConnector)

433 + connector.connector_scheduler = MagicMock()

434 + connector._kv_cache_events = None

435 + complete = _make_block_stored(b"complete", group_idx=0)

436 + pending = _make_block_stored(b"pending", group_idx=0)

437 +

438 + def poll(worker_events: list[BlockStored]) -> list[BlockStored]:

439 + events = mooncake_store_connector.MooncakeStoreKVEvents(

440 + num_workers=4,

441 + group_tp_replication_factors=(2,),

442 + )

443 + events.add_events(worker_events)

444 + connector.update_connector_output(KVConnectorOutput(kv_cache_events=events))

445 + return list(connector.take_events())

446 +

447 + assert poll([complete, pending]) == []

448 + assert poll([complete]) == [complete]

449 + assert poll([pending, pending]) == [pending]

450 + assert connector._kv_cache_events is None

451 +

452 +

453 # ============================================================

454 # reset_cache() — RL hard-reset path via typed LookupKey protocol

455 # ============================================================

tests/v1/kv_connector/unit/test_mooncake_store_layout.py new

+180

@@ -0,0 +1,180 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Tests for Mooncake Store payload layouts."""

4 +

5 +import ctypes

6 +import random

7 +

8 +import pytest

9 +import torch

10 +

11 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

12 + ChunkedTokenDatabase,

13 + KeyMetadata,

14 + LBHNCStoreLayout,

15 + LBNHCStoreLayout,

16 + RankLocalStoreLayout,

17 +)

18 +from vllm.utils.math_utils import cdiv

19 +

20 +BLOCK_SIZE = 128

21 +

22 +

23 +def _rank_local_layout(num_regions: int, num_block_lens: int) -> RankLocalStoreLayout:

24 + metadata = KeyMetadata("test-model", 1, 0, 0, 0)

25 + layout = RankLocalStoreLayout(metadata, BLOCK_SIZE, BLOCK_SIZE)

26 + layout.set_kv_caches_base_addr(

27 + [0x7F00_0000_0000 + i * (1 << 30) for i in range(num_regions)]

28 + )

29 + layout.set_block_len([30_208 + 512 * i for i in range(num_block_lens)])

30 + return layout

31 +

32 +

33 +def _rank_local_reference(

34 + layout: RankLocalStoreLayout, start: int, end: int, block_ids: list[int]

35 +) -> tuple[list[int], list[int], int]:

36 + block_id = block_ids[start // layout.block_size]

37 + length = len(layout.block_len)

38 + addrs = [

39 + base_addr + block_id * layout.block_len[index % length]

40 + for index, base_addr in enumerate(layout.kv_caches_base_addr)

41 + ]

42 + sizes = [

43 + layout.block_len[index % length] * cdiv(end - start, layout.block_size)

44 + for index in range(len(layout.kv_caches_base_addr))

45 + ]

46 + return addrs, sizes, block_id

47 +

48 +

49 +@pytest.mark.parametrize("num_regions,num_block_lens", [(96, 96), (96, 2), (1, 1)])

50 +def test_rank_local_descriptors_match_reference(num_regions: int, num_block_lens: int):

51 + layout = _rank_local_layout(num_regions, num_block_lens)

52 + rng = random.Random(0)

53 + block_ids = [rng.randrange(0, 1 << 20) for _ in range(300)]

54 + chunks = []

55 + block = 0

56 + while block < len(block_ids) - 4:

57 + span = rng.choice([1, 1, 1, 2, 4])

58 + chunks.append((block * BLOCK_SIZE, (block + span) * BLOCK_SIZE))

59 + block += span + rng.choice([0, 1])

60 +

61 + addrs, sizes, selected_blocks = layout.prepare_values(

62 + chunks, block_ids, [0] * len(chunks)

63 + )

64 +

65 + for chunk, chunk_addrs, chunk_sizes, block_id in zip(

66 + chunks, addrs, sizes, selected_blocks, strict=True

67 + ):

68 + assert (chunk_addrs, chunk_sizes, block_id) == _rank_local_reference(

69 + layout, *chunk, block_ids

70 + )

71 + assert all(type(addr) is int for addr in chunk_addrs)

72 + assert type(block_id) is int

73 +

74 +

75 +def test_rank_local_descriptors_handle_empty_and_invalid_chunks():

76 + layout = _rank_local_layout(4, 4)

77 + assert layout.prepare_values([], [1, 2, 3], []) == ([], [], [])

78 + with pytest.raises(AssertionError):

79 + layout.prepare_values([(0, BLOCK_SIZE + 1)], [0, 1], [0])

80 +

81 +

82 +@pytest.mark.parametrize(

83 + ("layout_cls", "store_format"),

84 + [

85 + (LBHNCStoreLayout, "tp_shared_lbhnc"),

86 + (LBNHCStoreLayout, "tp_shared_lbnhc"),

87 + ],

88 +)

89 +def test_tp_shared_layout_owns_store_namespace(layout_cls, store_format):

90 + assert (

91 + layout_cls.shared_namespace(4, 2)

92 + == f"@store_tp:4@store_pp:2@store_format:{store_format}"

93 + )

94 +

95 +

96 +def test_rank_local_database_api_rejects_tp_shared_layout():

97 + metadata = KeyMetadata("test-model", 0, 0, 0, 0)

98 + layout = LBHNCStoreLayout(

99 + metadata,

100 + block_size=16,

101 + hash_block_size=16,

102 + local_tp_size=2,

103 + store_tp_size=4,

104 + tp_rank=0,

105 + num_kv_heads=8,

106 + )

107 + database = ChunkedTokenDatabase(metadata, 16, store_layout=layout)

108 +

109 + with pytest.raises(RuntimeError, match="rank-local"):

110 + database.prepare_value_for_block(0)

111 +

112 +

113 +@pytest.mark.parametrize(

114 + ("layout_cls", "producer_strides", "consumer_strides"),

115 + [

116 + (LBHNCStoreLayout, (128, 64, 4, 1), (256, 64, 4, 1)),

117 + (LBNHCStoreLayout, (128, 4, 8, 1), (256, 4, 16, 1)),

118 + ],

119 +)

120 +def test_tp_shared_layout_round_trip_across_tp_sizes(

121 + layout_cls, producer_strides, consumer_strides

122 +):

123 + block_size = 16

124 + shape = (1, 2, block_size, 4)

125 + stored: dict[int, bytes] = {}

126 + producer_tensors = []

127 +

128 + for tp_rank in range(4):

129 + tensor = torch.empty_strided(shape, producer_strides, dtype=torch.float16)

130 + tensor.copy_(

131 + torch.arange(128, dtype=torch.float16).view(shape) + tp_rank * 1000

132 + )

133 + producer_tensors.append(tensor)

134 + metadata = KeyMetadata("test-model", tp_rank, 0, 0, 0)

135 + layout = layout_cls(

136 + metadata,

137 + block_size,

138 + block_size,

139 + local_tp_size=4,

140 + store_tp_size=4,

141 + tp_rank=tp_rank,

142 + num_kv_heads=8,

143 + )

144 + layout.register_kv_caches([tensor], 1)

145 + addrs, sizes, _ = layout.prepare_values([(0, block_size)], [0], [tp_rank])

146 + stored[tp_rank] = b"".join(

147 + ctypes.string_at(addr, size)

148 + for addr, size in zip(addrs[0], sizes[0], strict=True)

149 + )

150 +

151 + for tp_rank in range(2):

152 + tensor = torch.empty_strided(

153 + (1, 4, block_size, 4), consumer_strides, dtype=torch.float16

154 + )

155 + tensor.zero_()

156 + metadata = KeyMetadata("test-model", tp_rank, 0, 0, 0)

157 + layout = layout_cls(

158 + metadata,

159 + block_size,

160 + block_size,

161 + local_tp_size=2,

162 + store_tp_size=4,

163 + tp_rank=tp_rank,

164 + num_kv_heads=8,

165 + )

166 + layout.register_kv_caches([tensor], 1)

167 + shard_ids = layout.local_shard_ids

168 + addrs, sizes, _ = layout.prepare_values(

169 + [(0, block_size)] * len(shard_ids), [0], shard_ids

170 + )

171 + for shard_id, shard_addrs, shard_sizes in zip(

172 + shard_ids, addrs, sizes, strict=True

173 + ):

174 + offset = 0

175 + for addr, size in zip(shard_addrs, shard_sizes, strict=True):

176 + ctypes.memmove(addr, stored[shard_id][offset : offset + size], size)

177 + offset += size

178 +

179 + expected = torch.cat(producer_tensors[tp_rank * 2 : tp_rank * 2 + 2], dim=1)

180 + torch.testing.assert_close(tensor, expected)

tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py deleted

-90

@@ -1,90 +0,0 @@

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

+586 -16

@@ -17,6 +17,7 @@ import pytest

17 import torch

19 from tests.v1.attention.utils import dense_kv_cache_views

20 +from vllm.distributed.kv_events import KVEventAggregator

21 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (

22 rdma_utils,

23 )

@@ -30,9 +31,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

31 BlobBlockHashes,

32 ChunkedTokenDatabase,

33 KeyMetadata,

34 + LBHNCStoreLayout,

35 + LBNHCStoreLayout,

36 LoadSpec,

37 MooncakeLookupResult,

38 PoolKey,

39 + RankLocalStoreLayout,

40 ReqMeta,

41 TailKeyBoundary,

42 )

@@ -44,8 +48,19 @@ from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash

48 from vllm.v1.kv_cache_interface import (

49 FullAttentionSpec,

50 KVCacheGroupSpec,

47 - KVCacheLayout,

51 + RSWASpec,

52 )

53 +from vllm.v1.kv_cache_layout import KVCacheLayout

54 +

55 +_TP_SHARED_NAMESPACE = "@store_tp:4@store_pp:1@store_format:tp_shared_lbhnc"

56 +_TP_SHARED_LBNHC_NAMESPACE = "@store_tp:4@store_pp:1@store_format:tp_shared_lbnhc"

57 +

58 +

59 +def _tp_shared_prefix(tp_rank: int) -> str:

60 + return (

61 + f"test-model{_TP_SHARED_NAMESPACE}@tp_rank:{tp_rank}"

62 + "@pcp0@dcp0@pp_rank:0@group:0"

63 + )

66 class _RecordingBlockHashes:

@@ -78,6 +93,46 @@ def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:

93 )

96 +def _make_tp_shared_db(

97 + tp_rank: int = 0,

98 +) -> tuple[ChunkedTokenDatabase, torch.Tensor]:

99 + num_blocks = 2

100 + local_heads = 4

101 + block_size = 16

102 + content_size = 4

103 + strides = (

104 + local_heads * block_size * content_size,

105 + block_size * content_size,

106 + content_size,

107 + 1,

108 + )

109 + tensor = torch.empty_strided(

110 + (num_blocks, local_heads, block_size, content_size),

111 + strides,

112 + dtype=torch.float16,

113 + )

114 + metadata = KeyMetadata(

115 + "test-model",

116 + tp_rank,

117 + 0,

118 + 0,

119 + 0,

120 + store_namespace=_TP_SHARED_NAMESPACE,

121 + )

122 + layout = LBHNCStoreLayout(

123 + metadata,

124 + block_size,

125 + block_size,

126 + local_tp_size=2,

127 + store_tp_size=4,

128 + tp_rank=tp_rank,

129 + num_kv_heads=8,

130 + )

131 + db = ChunkedTokenDatabase(metadata, block_size, store_layout=layout)

132 + layout.register_kv_caches([tensor], num_blocks)

133 + return db, tensor

134 +

135 +

136 def _make_store_sending_thread(

137 store: MagicMock,

138 *,

@@ -243,11 +298,15 @@ def _make_vllm_config(

298 rank: int = 0,

299 decode_context_parallel_size: int = 1,

300 kv_role: str = "kv_both",

301 + pipeline_parallel_size: int = 1,

302 + kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC,

303 ) -> SimpleNamespace:

304 + cache_config = SimpleNamespace(block_size=16, num_gpu_blocks=10)

305 + cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout

306 return SimpleNamespace(

307 model_config=_FakeModelConfig(),

308 parallel_config=SimpleNamespace(

250 - pipeline_parallel_size=1,

309 + pipeline_parallel_size=pipeline_parallel_size,

310 rank=rank,

311 data_parallel_index=0,

312 decode_context_parallel_size=decode_context_parallel_size,

@@ -256,7 +315,7 @@ def _make_vllm_config(

315 kv_transfer_config=_FakeKVTransferConfig(

316 kv_role=kv_role, extra_config=extra_config

317 ),

259 - cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10),

318 + cache_config=cache_config,

319 kv_events_config=SimpleNamespace(enable_kv_cache_events=False),

320 speculative_config=None,

321 )

@@ -334,6 +393,138 @@ def test_pool_key_to_string_without_prefix_is_unchanged():

393 )

394

395

396 +def test_pool_key_store_tp_namespace():

397 + key = PoolKey(

398 + KeyMetadata(

399 + "test-model",

400 + 2,

401 + 0,

402 + 0,

403 + 0,

404 + store_namespace=_TP_SHARED_NAMESPACE,

405 + ),

406 + "deadbeef",

407 + )

408 +

409 + assert key.to_string() == f"{_tp_shared_prefix(2)}@deadbeef"

410 +

411 +

412 +def test_tp_shared_sending_writes_each_store_shard():

413 + store = MagicMock()

414 + store.batch_is_exist.return_value = [0, 0, 0, 0]

415 + store.batch_put_from_multi_buffers.return_value = [1, 1, 1, 1]

416 + db, _ = _make_tp_shared_db()

417 + thread = _make_store_sending_thread(

418 + store,

419 + token_databases=[db],

420 + replicate_config=SimpleNamespace(),

421 + )

422 + req = _make_store_req("req", [b"h0", b"h1"])

423 + _run_store_req(thread, req)

424 +

425 + keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args

426 + assert keys == [

427 + f"{_tp_shared_prefix(0)}@6830",

428 + f"{_tp_shared_prefix(1)}@6830",

429 + f"{_tp_shared_prefix(0)}@6831",

430 + f"{_tp_shared_prefix(1)}@6831",

431 + ]

432 + assert len(addrs) == len(sizes) == 4

433 +

434 +

435 +def test_tp_shared_kv_events_aggregate_once_per_logical_block():

436 + aggregator = KVEventAggregator(num_workers=2)

437 +

438 + for tp_rank in range(2):

439 + store = MagicMock()

440 + store.batch_is_exist.return_value = [0, 0]

441 + store.batch_put_from_multi_buffers.return_value = [256, 256]

442 + db, _ = _make_tp_shared_db(tp_rank)

443 + thread = _make_store_sending_thread(

444 + store,

445 + token_databases=[db],

446 + tp_rank=tp_rank,

447 + replicate_config=SimpleNamespace(),

448 + )

449 + thread.enable_kv_event = True

450 +

451 + _run_store_req(thread, _make_event_store_req(16))

452 +

453 + events = thread.get_kv_events()

454 + assert len(events) == 1

455 + aggregator.add_events(events)

456 +

457 + assert len(aggregator.get_common_events()) == 1

458 +

459 +

460 +def test_tp_shared_kv_event_waits_for_all_missing_shards():

461 + store = MagicMock()

462 + store.batch_is_exist.return_value = [0, 0]

463 + store.batch_put_from_multi_buffers.return_value = [256, -1]

464 + db, _ = _make_tp_shared_db()

465 + thread = _make_store_sending_thread(

466 + store,

467 + token_databases=[db],

468 + replicate_config=SimpleNamespace(),

469 + )

470 + thread.enable_kv_event = True

471 +

472 + _run_store_req(thread, _make_event_store_req(16))

473 +

474 + assert thread.get_kv_events() == []

475 +

476 +

477 +def test_tp_shared_kv_event_counts_existing_shards_as_satisfied():

478 + store = MagicMock()

479 + store.batch_is_exist.return_value = [1, 0]

480 + store.batch_put_from_multi_buffers.return_value = [256]

481 + db, _ = _make_tp_shared_db()

482 + thread = _make_store_sending_thread(

483 + store,

484 + token_databases=[db],

485 + replicate_config=SimpleNamespace(),

486 + )

487 + thread.enable_kv_event = True

488 +

489 + _run_store_req(thread, _make_event_store_req(16))

490 +

491 + assert len(thread.get_kv_events()) == 1

492 +

493 +

494 +def test_tp_shared_receiving_reads_each_local_store_shard():

495 + from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec

496 +

497 + store = MagicMock()

498 + store.batch_get_into_multi_buffers.return_value = [1, 1, 1, 1]

499 + db, _ = _make_tp_shared_db()

500 + spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

501 + coord = mooncake_store_worker.MooncakeStoreCoordinator(

502 + [KVCacheGroupSpec(["layer0"], spec)],

503 + scheduler_block_size=16,

504 + hash_block_size=16,

505 + )

506 + thread = mooncake_store_worker.KVCacheStoreRecvingThread(

507 + store=store,

508 + token_databases=[db],

509 + block_size=16,

510 + tp_rank=0,

511 + ready_event=threading.Event(),

512 + coord=coord,

513 + )

514 + thread.request_queue.task_done = MagicMock()

515 +

516 + thread._handle_request(_make_load_req("req", [b"h0", b"h1"], token_len=32))

517 +

518 + keys, addrs, sizes = store.batch_get_into_multi_buffers.call_args.args

519 + assert keys == [

520 + f"{_tp_shared_prefix(0)}@6830",

521 + f"{_tp_shared_prefix(1)}@6830",

522 + f"{_tp_shared_prefix(0)}@6831",

523 + f"{_tp_shared_prefix(1)}@6831",

524 + ]

525 + assert len(addrs) == len(sizes) == 4

526 +

527 +

528 def test_pool_key_cache_prefix_namespaces_and_disambiguates():

529 """A non-empty cache_prefix is prepended, and two instances with

530 different prefixes never collide on identical block hashes."""

@@ -1155,8 +1346,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():

1346 )

1347 db0.set_kv_caches_base_addr([0x1000])

1348 db0.set_block_len([256])

1158 - db0.prepare_values = MagicMock(wraps=db0.prepare_values)

1159 - db0.prepare_value = MagicMock(side_effect=AssertionError("scalar path called"))

1349 + db0.store_layout.prepare_values = MagicMock(wraps=db0.store_layout.prepare_values)

1350

1351 db1 = ChunkedTokenDatabase(

1352 KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),

@@ -1164,8 +1354,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():

1354 )

1355 db1.set_kv_caches_base_addr([0x2000])

1356 db1.set_block_len([512])

1167 - db1.prepare_values = MagicMock(wraps=db1.prepare_values)

1168 - db1.prepare_value = MagicMock(side_effect=AssertionError("scalar path called"))

1357 + db1.store_layout.prepare_values = MagicMock(wraps=db1.store_layout.prepare_values)

1358

1359 thread = _make_store_sending_thread(

1360 store,

@@ -1183,10 +1372,12 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():

1372 ),

1373 )

1374

1186 - db0.prepare_value.assert_not_called()

1187 - db1.prepare_value.assert_not_called()

1188 - db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [1, 2, 3])

1189 - db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [3, 2, 1])

1375 + db0.store_layout.prepare_values.assert_called_once_with(

1376 + [(0, 16), (32, 48)], [1, 2, 3], [0, 0]

1377 + )

1378 + db1.store_layout.prepare_values.assert_called_once_with(

1379 + [(16, 32), (32, 48)], [3, 2, 1], [0, 0]

1380 + )

1381

1382 keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args

1383 assert [key.rsplit("@", 1)[-1] for key in keys] == [

@@ -1836,8 +2027,10 @@ def test_recv_thread_splits_disk_offload_loads_by_budget():

2027 assert second_keys == [

2028 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",

2029 ]

1839 - base_addr = thread.token_databases[0].kv_caches_base_addr[0]

1840 - block_len = thread.token_databases[0].block_len[0]

2030 + layout = thread.token_databases[0].store_layout

2031 + assert isinstance(layout, RankLocalStoreLayout)

2032 + base_addr = layout.kv_caches_base_addr[0]

2033 + block_len = layout.block_len[0]

2034 assert first_addrs == [[base_addr], [base_addr + block_len]]

2035 assert second_addrs == [[base_addr + 2 * block_len]]

2036 expected_size = block_len

@@ -2004,6 +2197,294 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled(

2197 assert w.disk_offload_buffer_budget_bytes is None

2198

2199

2200 +def test_worker_enables_lbhnc_store_tp_layout(tmp_path, monkeypatch):

2201 + store = MagicMock()

2202 + store.setup.return_value = 0

2203 + _install_fake_mooncake(monkeypatch, store)

2204 + _patch_worker_runtime(monkeypatch, tp_size=2)

2205 + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)

2206 + monkeypatch.setenv(

2207 + "MOONCAKE_CONFIG_PATH",

2208 + _write_mooncake_config(

2209 + tmp_path,

2210 + {

2211 + "metadata_server": "http://metadata/endpoint",

2212 + "protocol": "tcp",

2213 + "master_server_address": "10.0.0.7:50051",

2214 + },

2215 + ),

2216 + )

2217 +

2218 + w = worker.MooncakeStoreWorker(

2219 + _make_vllm_config(extra_config={"store_tp_size": 4}),

2220 + _make_kv_cache_config(),

2221 + )

2222 +

2223 + assert w.store_tp_size == 4

2224 + assert w.token_dbs[0].metadata.store_namespace == _TP_SHARED_NAMESPACE

2225 + assert isinstance(w.token_dbs[0].store_layout, LBHNCStoreLayout)

2226 + assert w.token_dbs[0].store_layout.local_shard_ids == (0, 1)

2227 + assert len(w._lookup_key_prefixes[0]) == 4

2228 +

2229 +

2230 +def test_worker_enables_lbnhc_store_tp_layout(tmp_path, monkeypatch):

2231 + store = MagicMock()

2232 + store.setup.return_value = 0

2233 + _install_fake_mooncake(monkeypatch, store)

2234 + _patch_worker_runtime(monkeypatch, tp_size=2)

2235 + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)

2236 + monkeypatch.setenv(

2237 + "MOONCAKE_CONFIG_PATH",

2238 + _write_mooncake_config(

2239 + tmp_path,

2240 + {

2241 + "metadata_server": "http://metadata/endpoint",

2242 + "protocol": "tcp",

2243 + "master_server_address": "10.0.0.7:50051",

2244 + },

2245 + ),

2246 + )

2247 +

2248 + with patch.object(worker.logger, "warning_once") as warning_once:

2249 + w = worker.MooncakeStoreWorker(

2250 + _make_vllm_config(

2251 + extra_config={"store_tp_size": 4},

2252 + kv_cache_layout=KVCacheLayout.LBNHC,

2253 + ),

2254 + _make_kv_cache_config(),

2255 + )

2256 +

2257 + warning_once.assert_called_once()

2258 + assert w.store_tp_size == 4

2259 + assert w.token_dbs[0].metadata.store_namespace == _TP_SHARED_LBNHC_NAMESPACE

2260 + assert isinstance(w.token_dbs[0].store_layout, LBNHCStoreLayout)

2261 + assert w.token_dbs[0].store_layout.local_shard_ids == (0, 1)

2262 + assert len(w._lookup_key_prefixes[0]) == 4

2263 +

2264 +

2265 +def test_tp_sharded_layout_rejects_full_attention_subclasses():

2266 + spec = RSWASpec(

2267 + block_size=16,

2268 + num_kv_heads=8,

2269 + head_size=64,

2270 + dtype=None,

2271 + rswa_window=32,

2272 + )

2273 + store_worker = object.__new__(worker.MooncakeStoreWorker)

2274 + store_worker.pcp_size = 1

2275 + store_worker.dcp_size = 1

2276 + store_worker._kv_cache_groups = [SimpleNamespace(kv_cache_spec=spec)]

2277 +

2278 + assert not store_worker._supports_tp_sharded_store_layout(LBHNCStoreLayout, {})

2279 +

2280 +

2281 +@pytest.mark.parametrize(

2282 + ("extra_config", "expected"),

2283 + [

2284 + ({"store_tp_size": 4}, 4),

2285 + (

2286 + {

2287 + "enable_store_tp_lcm": True,

2288 + "prefill_tp_sizes": [3, 4, 6],

2289 + },

2290 + 12,

2291 + ),

2292 + (

2293 + {

2294 + "enable_store_tp_lcm": False,

2295 + "prefill_tp_sizes": [4, 2],

2296 + },

2297 + None,

2298 + ),

2299 + ({"enable_store_tp_lcm": True, "prefill_tp_sizes": []}, None),

2300 + ({"enable_store_tp_lcm": True, "prefill_tp_sizes": [4, 0]}, None),

2301 + ({"enable_store_tp_lcm": True, "prefill_tp_sizes": [4, 2.0]}, None),

2302 + ({"store_tp_size": 0}, None),

2303 + ({"store_tp_size": -1}, None),

2304 + ({"store_tp_size": True}, None),

2305 + ({"store_tp_size": 4.0}, None),

2306 + ],

2307 +)

2308 +def test_resolve_store_tp_size(extra_config: dict[str, object], expected: int | None):

2309 + assert worker.resolve_store_tp_size(extra_config) == expected

2310 +

2311 +

2312 +def test_lcm_store_tp_gives_prefill_and_decode_common_namespace(tmp_path, monkeypatch):

2313 + store = MagicMock()

2314 + store.setup.return_value = 0

2315 + _install_fake_mooncake(monkeypatch, store)

2316 + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)

2317 + monkeypatch.setenv(

2318 + "MOONCAKE_CONFIG_PATH",

2319 + _write_mooncake_config(

2320 + tmp_path,

2321 + {

2322 + "metadata_server": "http://metadata/endpoint",

2323 + "protocol": "tcp",

2324 + "master_server_address": "10.0.0.7:50051",

2325 + },

2326 + ),

2327 + )

2328 + lcm_config = {

2329 + "enable_store_tp_lcm": True,

2330 + "prefill_tp_sizes": [4, 2],

2331 + }

2332 +

2333 + workers = []

2334 + topologies = [

2335 + (4, 2, "kv_both", (2,)),

2336 + (2, 1, "kv_both", (2, 3)),

2337 + (2, 1, "kv_consumer", (2, 3)),

2338 + ]

2339 + for tp_size, tp_rank, kv_role, expected_shards in topologies:

2340 + _patch_worker_runtime(

2341 + monkeypatch,

2342 + tp_size=tp_size,

2343 + tp_rank=tp_rank,

2344 + )

2345 + extra_config = dict(lcm_config)

2346 + if kv_role == "kv_consumer":

2347 + extra_config["save_decode_cache"] = True

2348 + store_worker = worker.MooncakeStoreWorker(

2349 + _make_vllm_config(extra_config=extra_config, kv_role=kv_role),

2350 + _make_kv_cache_config(),

2351 + )

2352 + assert store_worker.store_tp_size == 4

2353 + assert isinstance(store_worker.token_dbs[0].store_layout, LBHNCStoreLayout)

2354 + assert store_worker.token_dbs[0].store_layout.local_shard_ids == expected_shards

2355 + workers.append(store_worker)

2356 +

2357 + shard_keys = [

2358 + store_worker.token_dbs[0].store_layout.key_for(2, BlockHash(b"h"))

2359 + for store_worker in workers

2360 + ]

2361 + assert len(set(shard_keys)) == 1

2362 + assert workers[-1].can_put

2363 +

2364 +

2365 +@pytest.mark.parametrize(

2366 + ("tp_size", "prefill_tp_sizes"),

2367 + [

2368 + (3, [4, 2]),

2369 + (2, [4, 3]),

2370 + (2, []),

2371 + (2, [4, 0]),

2372 + ],

2373 +)

2374 +def test_lcm_store_tp_falls_back_when_topology_is_incompatible(

2375 + tmp_path,

2376 + monkeypatch,

2377 + tp_size: int,

2378 + prefill_tp_sizes: list[int],

2379 +):

2380 + store = MagicMock()

2381 + store.setup.return_value = 0

2382 + _install_fake_mooncake(monkeypatch, store)

2383 + _patch_worker_runtime(monkeypatch, tp_size=tp_size)

2384 + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)

2385 + monkeypatch.setenv(

2386 + "MOONCAKE_CONFIG_PATH",

2387 + _write_mooncake_config(

2388 + tmp_path,

2389 + {

2390 + "metadata_server": "http://metadata/endpoint",

2391 + "protocol": "tcp",

2392 + "master_server_address": "10.0.0.7:50051",

2393 + },

2394 + ),

2395 + )

2396 +

2397 + store_worker = worker.MooncakeStoreWorker(

2398 + _make_vllm_config(

2399 + extra_config={

2400 + "enable_store_tp_lcm": True,

2401 + "prefill_tp_sizes": prefill_tp_sizes,

2402 + }

2403 + ),

2404 + _make_kv_cache_config(),

2405 + )

2406 +

2407 + assert store_worker.store_tp_size is None

2408 + assert type(store_worker.token_dbs[0]) is ChunkedTokenDatabase

2409 + assert "@store_format:rank_local_tp" in store_worker.token_dbs[0].key_for(

2410 + BlockHash(b"h")

2411 + )

2412 +

2413 +

2414 +@pytest.mark.parametrize("store_tp_size", [1, 3, 6])

2415 +def test_worker_keeps_rank_local_layout_for_unsupported_store_tp(

2416 + tmp_path,

2417 + monkeypatch,

2418 + store_tp_size: int,

2419 +):

2420 + store = MagicMock()

2421 + store.setup.return_value = 0

2422 + _install_fake_mooncake(monkeypatch, store)

2423 + _patch_worker_runtime(monkeypatch, tp_size=2)

2424 + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)

2425 + monkeypatch.setenv(

2426 + "MOONCAKE_CONFIG_PATH",

2427 + _write_mooncake_config(

2428 + tmp_path,

2429 + {

2430 + "metadata_server": "http://metadata/endpoint",

2431 + "protocol": "tcp",

2432 + "master_server_address": "10.0.0.7:50051",

2433 + },

2434 + ),

2435 + )

2436 +

2437 + w = worker.MooncakeStoreWorker(

2438 + _make_vllm_config(

2439 + extra_config={"store_tp_size": store_tp_size},

2440 + kv_cache_layout=KVCacheLayout.LBNHC,

2441 + ),

2442 + _make_kv_cache_config(),

2443 + )

2444 +

2445 + assert w.store_tp_size is None

2446 + assert type(w.token_dbs[0]) is ChunkedTokenDatabase

2447 + key = w.token_dbs[0].key_for(BlockHash(b"h"))

2448 + assert "@store_tp:" not in key

2449 + assert "@store_pp:1" in key

2450 + assert "@store_format:rank_local_tp2_layout_LBNHC" in key

2451 +

2452 +

2453 +def test_store_tp_opt_in_isolates_different_pp_sizes(tmp_path, monkeypatch):

2454 + store = MagicMock()

2455 + store.setup.return_value = 0

2456 + _install_fake_mooncake(monkeypatch, store)

2457 + _patch_worker_runtime(monkeypatch, tp_size=2)

2458 + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)

2459 + monkeypatch.setenv(

2460 + "MOONCAKE_CONFIG_PATH",

2461 + _write_mooncake_config(

2462 + tmp_path,

2463 + {

2464 + "metadata_server": "http://metadata/endpoint",

2465 + "protocol": "tcp",

2466 + "master_server_address": "10.0.0.7:50051",

2467 + },

2468 + ),

2469 + )

2470 +

2471 + keys = []

2472 + for pp_size in (1, 2):

2473 + w = worker.MooncakeStoreWorker(

2474 + _make_vllm_config(

2475 + extra_config={"store_tp_size": 4},

2476 + pipeline_parallel_size=pp_size,

2477 + ),

2478 + _make_kv_cache_config(),

2479 + )

2480 + layout = w.token_dbs[0].store_layout

2481 + keys.append(layout.key_for(layout.local_shard_ids[0], BlockHash(b"h")))

2482 +

2483 + assert "@store_pp:1" in keys[0]

2484 + assert "@store_pp:2" in keys[1]

2485 + assert keys[0] != keys[1]

2486 +

2487 +

2488 def test_save_decode_cache_keeps_transfer_path_enabled(tmp_path, monkeypatch):

2489 store = MagicMock()

2490 store.setup.return_value = 0

@@ -2185,6 +2666,58 @@ def test_worker_put_striding_covers_every_rank_get_namespace(

2666 )

2667

2668

2669 +def test_mqa_p4_to_d2_uses_shared_rank_zero_namespace(tmp_path, monkeypatch):

2670 + store = MagicMock()

2671 + store.setup.return_value = 0

2672 + _install_fake_mooncake(monkeypatch, store)

2673 + monkeypatch.setenv(

2674 + "MOONCAKE_CONFIG_PATH",

2675 + _write_mooncake_config(

2676 + tmp_path,

2677 + {

2678 + "metadata_server": "http://metadata/endpoint",

2679 + "protocol": "tcp",

2680 + "master_server_address": "10.0.0.7:50051",

2681 + },

2682 + ),

2683 + )

2684 +

2685 + block_hash = BlockHash(b"hash")

2686 + put_keys: set[str] = set()

2687 + get_keys: set[str] = set()

2688 + for tp_size, target in ((4, put_keys), (2, get_keys)):

2689 + for tp_rank in range(tp_size):

2690 + _patch_worker_runtime(monkeypatch, tp_rank=tp_rank, tp_size=tp_size)

2691 + w = worker.MooncakeStoreWorker(

2692 + _make_vllm_config(

2693 + rank=tp_rank,

2694 + extra_config={"store_tp_size": 4},

2695 + ),

2696 + _make_kv_cache_config(),

2697 + )

2698 +

2699 + assert w.store_tp_size is None

2700 + assert w.token_dbs[0].metadata.tp_rank == 0

2701 + key = w.token_dbs[0].key_for(block_hash)

2702 + if tp_size == 4:

2703 + put_step = w._group_tp_replication_factors[0]

2704 + if tp_rank % put_step == 0:

2705 + target.add(key)

2706 + else:

2707 + target.add(key)

2708 +

2709 + assert (

2710 + put_keys

2711 + == get_keys

2712 + == {

2713 + (

2714 + "test-model@store_pp:1@store_format:tp_shared_mqa"

2715 + "@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@68617368"

2716 + )

2717 + }

2718 + )

2719 +

2720 +

2721 def test_requester_worker_group_semantics_falls_back_without_group_ids(

2722 tmp_path,

2723 monkeypatch,

@@ -2396,9 +2929,11 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():

2929

2930 store = MagicMock()

2931 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

2399 - store.batch_put_from_multi_buffers.side_effect = (

2400 - lambda keys, addrs, sizes, *_args: ([256] * len(keys))

2401 - )

2932 +

2933 + def successful_put(keys, addrs, sizes, *_args):

2934 + return [256] * len(keys)

2935 +

2936 + store.batch_put_from_multi_buffers.side_effect = successful_put

2937

2938 full_spec = FullAttentionSpec(

2939 block_size=32, num_kv_heads=8, head_size=64, dtype=None

@@ -2754,6 +3289,7 @@ def _make_bare_worker(

3289 worker.recv_request_queue = queue.Queue()

3290 worker.finished_store_req = set()

3291 worker.tp_size = 1

3292 + worker.store_tp_size = None

3293 worker.num_kv_head = 1

3294 worker.pp_size = 1

3295 # Minimal single-full-attention-group config so the coordinator-based

@@ -2812,6 +3348,40 @@ def test_lookup_key_prefixes_cover_dcp_rank_namespaces():

3348 )

3349

3350

3351 +def test_lookup_key_prefixes_cover_store_tp_shards():

3352 + worker = _make_bare_worker()

3353 + worker.store_tp_size = 4

3354 + metadata = KeyMetadata(

3355 + "test-model",

3356 + 0,

3357 + 0,

3358 + 0,

3359 + 0,

3360 + group_id=0,

3361 + store_namespace=_TP_SHARED_NAMESPACE,

3362 + )

3363 + worker.token_dbs = [

3364 + ChunkedTokenDatabase(

3365 + metadata,

3366 + block_size=16,

3367 + store_layout=LBHNCStoreLayout(

3368 + metadata,

3369 + block_size=16,

3370 + hash_block_size=16,

3371 + local_tp_size=2,

3372 + store_tp_size=4,

3373 + tp_rank=0,

3374 + num_kv_heads=8,

3375 + ),

3376 + )

3377 + ]

3378 + worker._init_lookup_key_prefixes()

3379 +

3380 + assert worker._lookup_key_prefixes[0] == (

3381 + *(_tp_shared_prefix(tp_rank) for tp_rank in range(4)),

3382 + )

3383 +

3384 +

3385 def test_lookup_key_prefixes_cover_pcp_rank_namespaces():

3386 worker = _make_bare_worker()

3387 worker.tp_size = 4

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

+66 -25

@@ -11,16 +11,17 @@ and consumer instances read/write KV to/from the store independently,

11 enabling prefix caching via hash-based deduplication.

12 """

14 -from collections.abc import Iterable

14 +from collections import Counter

15 +from collections.abc import Iterable, Sequence

16 from typing import Any

18 import torch

20 from vllm.config import VllmConfig

21 from vllm.distributed.kv_events import (

22 + BlockStored,

23 KVCacheEvent,

24 KVConnectorKVEvents,

23 - KVEventAggregator,

25 )

26 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

27 KVConnectorBase_V1,

@@ -56,31 +57,66 @@ logger = init_logger(__name__)

57 class MooncakeStoreKVEvents(KVConnectorKVEvents):

58 """KV event aggregation for MooncakeStoreConnector."""

59 - def __init__(self, num_workers: int) -> None:

60 - self._aggregator = KVEventAggregator(num_workers)

60 + def __init__(

61 + self,

62 + num_workers: int,

63 + group_tp_replication_factors: Sequence[int] = (1,),

64 + ) -> None:

65 + if num_workers <= 0:

66 + raise ValueError("num_workers must be greater than zero.")

67 + if any(factor <= 0 for factor in group_tp_replication_factors):

68 + raise ValueError("TP replication factors must be greater than zero.")

69 + self._event_counter: Counter[KVCacheEvent] = Counter()

70 + self._num_workers = num_workers

71 + self._group_tp_replication_factors = tuple(group_tp_replication_factors)

73 def add_events(self, events: list[KVCacheEvent]) -> None:

63 - self._aggregator.add_events(events)

74 + if not isinstance(events, list):

75 + raise TypeError("events must be a list of KVCacheEvent.")

76 + self._event_counter.update(events)

77 +

78 + def _replication_factor(self, event: KVCacheEvent) -> int:

79 + if not isinstance(event, BlockStored) or event.group_idx is None:

80 + return 1

81 + return self._group_tp_replication_factors[event.group_idx]

82 +

83 + def _is_common_event(self, event: KVCacheEvent, count: int) -> bool:

84 + return count * self._replication_factor(event) >= self._num_workers

86 def aggregate(self) -> "MooncakeStoreKVEvents":

66 - common_events = self._aggregator.get_common_events()

67 - self._aggregator.clear_events()

68 - self._aggregator.add_events(common_events)

69 - self._aggregator.reset_workers()

87 + common_events = self.pop_common_events()

88 + self._event_counter.clear()

89 + self._event_counter.update(common_events)

90 + self._num_workers = 1

91 return self

93 + def pop_common_events(self) -> list[KVCacheEvent]:

94 + common_events = [

95 + event

96 + for event, count in self._event_counter.items()

97 + if self._is_common_event(event, count)

98 + ]

99 + for event in common_events:

100 + del self._event_counter[event]

101 + return common_events

102 +

103 + def has_events(self) -> bool:

104 + return bool(self._event_counter)

105 +

106 def increment_workers(self, count: int = 1) -> None:

73 - self._aggregator.increment_workers(count)

107 + if count <= 0:

108 + raise ValueError("count must be positive.")

109 + self._num_workers += count

110

111 def get_all_events(self) -> list[KVCacheEvent]:

76 - return self._aggregator.get_all_events()

112 + return list(self._event_counter.elements())

113

114 def get_number_of_workers(self) -> int:

79 - return self._aggregator.get_number_of_workers()

115 + return self._num_workers

116

117 def clear_events(self) -> None:

82 - self._aggregator.clear_events()

83 - self._aggregator.reset_workers()

118 + self._event_counter.clear()

119 + self._num_workers = 1

120

121 def __repr__(self) -> str:

122 return f"<MooncakeStoreKVEvents events={self.get_all_events()}>"

@@ -260,16 +296,13 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

296 self._kv_cache_events = kv_cache_events

297 else:

298 self._kv_cache_events.add_events(kv_cache_events.get_all_events())

263 - self._kv_cache_events.increment_workers(

264 - kv_cache_events.get_number_of_workers()

265 - )

299

300 def take_events(self) -> Iterable[KVCacheEvent]:

301 if self._kv_cache_events is not None:

269 - self._kv_cache_events.aggregate()

270 - yield from self._kv_cache_events.get_all_events()

271 - self._kv_cache_events.clear_events()

272 - self._kv_cache_events = None

302 + events = self._kv_cache_events.pop_common_events()

303 + if not self._kv_cache_events.has_events():

304 + self._kv_cache_events = None

305 + yield from events

306

307 # ============================================================

308 # Worker-side methods

@@ -321,11 +354,19 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

354 self,

355 ) -> MooncakeStoreKVEvents | None:

356 assert self.connector_worker is not None

324 - events = self.connector_worker.get_kv_events()

325 - if not events:

357 + if (

358 + not self.connector_worker.enable_kv_events

359 + or self.connector_worker.kv_send_thread is None

360 + ):

361 return None

327 -

328 - kv_events = MooncakeStoreKVEvents(num_workers=1)

362 + events = self.connector_worker.get_kv_events()

363 + # Empty containers still count this worker toward the poll's quorum.

364 + kv_events = MooncakeStoreKVEvents(

365 + num_workers=1,

366 + group_tp_replication_factors=(

367 + self.connector_worker.group_tp_replication_factors

368 + ),

369 + )

370 kv_events.add_events(events)

371 return kv_events

372

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

+392 -41

@@ -18,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (

18 )

19 from vllm.logger import init_logger

20 from vllm.utils.math_utils import cdiv

21 +from vllm.utils.torch_utils import is_non_overlapping_and_dense

22 from vllm.v1.core.kv_cache_utils import (

23 BlockHash,

24 BlockHashListWithBlockSize,

@@ -111,6 +112,9 @@ class KeyMetadata:

112 # share one Mooncake master without colliding on identical block hashes.

113 # Empty (the default) keeps keys byte-identical to the unprefixed format.

114 cache_prefix: str = ""

115 + # Complete namespace for an opt-in Store payload format. Empty keeps

116 + # historical keys byte-identical.

117 + store_namespace: str = ""

118

119

120 @dataclass(order=True)

@@ -125,6 +129,7 @@ class PoolKey:

129 (

130 self.key_metadata.cache_prefix,

131 self.key_metadata.model_name,

132 + self.key_metadata.store_namespace,

133 self.key_metadata.tp_rank,

134 self.key_metadata.pcp_rank,

135 self.key_metadata.dcp_rank,

@@ -148,6 +153,7 @@ class PoolKey:

153 return (

154 f"{prefix}"

155 f"{key_metadata.model_name}"

156 + f"{key_metadata.store_namespace}"

157 f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}"

158 f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}"

159 f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}"

@@ -165,61 +171,149 @@ class PoolKey:

171 )

172

173

168 -class ChunkedTokenDatabase:

169 - """Maps token positions to store keys and GPU memory addresses."""

174 +StoreShardId = int

175 +# Physical ranks encoded in a rank-local key: (TP, PCP, DCP, PP).

176 +RankNamespace = tuple[int, int, int, int]

177 +

178 +

179 +class StoreLayout:

180 + """Store payload layout for one KV-cache group on one local TP rank."""

181

182 def __init__(

183 self,

184 metadata: KeyMetadata,

185 block_size: int,

175 - hash_block_size: int | None = None,

176 - ):

186 + hash_block_size: int,

187 + ) -> None:

188 self.metadata = metadata

189 self.block_size = block_size

179 - self.hash_block_size = hash_block_size or block_size

180 - if self.block_size % self.hash_block_size != 0:

181 - raise ValueError(

182 - f"block_size ({self.block_size}) must be a multiple of "

183 - f"hash_block_size ({self.hash_block_size})"

184 - )

190 + self.hash_block_size = hash_block_size

191 +

192 + @property

193 + def local_shard_ids(self) -> tuple[StoreShardId, ...]:

194 + """Store objects contributed or consumed by this local rank."""

195 + raise NotImplementedError

196 +

197 + def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:

198 + raise NotImplementedError

199 +

200 + def lookup_key_prefixes(

201 + self,

202 + rank_namespaces: Sequence[RankNamespace],

203 + ) -> tuple[str, ...]:

204 + """Prefixes that must exist for a logical block to be reusable."""

205 + raise NotImplementedError

206 +

207 + def register_kv_caches(

208 + self,

209 + kv_caches: Sequence[torch.Tensor],

210 + num_blocks: int,

211 + ) -> None:

212 + """Build descriptors from the local KV cache tensors."""

213 + raise NotImplementedError

214 +

215 + def prepare_values(

216 + self,

217 + chunks: Sequence[tuple[int, int]],

218 + block_ids: list[int],

219 + shard_ids: Sequence[StoreShardId],

220 + ) -> tuple[list[list[int]], list[list[int]], list[int]]:

221 + """Build Store multi-buffer descriptors for logical chunks."""

222 + raise NotImplementedError

223 +

224 +

225 +class RankLocalStoreLayout(StoreLayout):

226 + """Historical rank-local Store payload layout."""

227 +

228 + def __init__(

229 + self,

230 + metadata: KeyMetadata,

231 + block_size: int,

232 + hash_block_size: int,

233 + ) -> None:

234 + super().__init__(metadata, block_size, hash_block_size)

235 + self._key_prefix = PoolKey.build_prefix(metadata)

236 self.kv_caches_base_addr: list[int] = []

237 self.block_len: list[int] = []

187 - self._key_prefix = PoolKey.build_prefix(metadata)

238

189 - def key_for(self, chunk_hash: BlockHash) -> str:

239 + @property

240 + def local_shard_ids(self) -> tuple[StoreShardId, ...]:

241 + return (0,)

242 +

243 + def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:

244 + assert shard_id == 0

245 return PoolKey.build_key_string(self._key_prefix, chunk_hash.hex())

246

192 - def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]):

193 - self.kv_caches_base_addr = kv_caches_base_addr

247 + def lookup_key_prefixes(

248 + self,

249 + rank_namespaces: Sequence[RankNamespace],

250 + ) -> tuple[str, ...]:

251 + return tuple(

252 + PoolKey.build_prefix(

253 + self.metadata,

254 + tp_rank=tp_rank,

255 + pcp_rank=pcp_rank,

256 + dcp_rank=dcp_rank,

257 + pp_rank=pp_rank,

258 + )

259 + for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces

260 + )

261

195 - def set_block_len(self, block_len: list[int]):

196 - self.block_len = block_len

262 + def set_kv_caches_base_addr(self, base_addrs: list[int]) -> None:

263 + self.kv_caches_base_addr = base_addrs

264

198 - def prepare_value(

199 - self, start: int, end: int, block_ids: list[int]

200 - ) -> tuple[list[int], list[int], int]:

201 - """Compute memory addresses and sizes for a single token range.

265 + def set_block_len(self, block_lens: list[int]) -> None:

266 + self.block_len = block_lens

267

203 - Returns:

204 - (addr_list, size_list, block_id)

205 - """

206 - addr_lists, size_lists, chunk_block_ids = self.prepare_values(

207 - ((start, end),), block_ids

208 - )

209 - return addr_lists[0], size_lists[0], chunk_block_ids[0]

268 + def register_kv_caches(

269 + self,

270 + kv_caches: Sequence[torch.Tensor],

271 + num_blocks: int,

272 + ) -> None:

273 + seen_region_ptrs: set[int] = set()

274 + base_addrs: list[int] = []

275 + block_lens: list[int] = []

276 + for cache in kv_caches:

277 + storage = cache.untyped_storage()

278 + storage_addr = storage.data_ptr()

279 + storage_len = storage.nbytes()

280 + if not is_non_overlapping_and_dense(cache[0]):

281 + for head_idx in range(cache.shape[1]):

282 + head_cache = cache[:, head_idx]

283 + assert is_non_overlapping_and_dense(head_cache[0])

284 + region_addr = head_cache.data_ptr()

285 + if region_addr in seen_region_ptrs:

286 + continue

287 + seen_region_ptrs.add(region_addr)

288 + base_addrs.append(region_addr)

289 + block_lens.append(head_cache.stride(0) * head_cache.element_size())

290 + elif cache.stride(0) * cache.element_size() * num_blocks == storage_len:

291 + if storage_addr in seen_region_ptrs:

292 + continue

293 + seen_region_ptrs.add(storage_addr)

294 + base_addrs.append(storage_addr)

295 + block_lens.append(storage_len // num_blocks)

296 + else:

297 + region_addr = cache.data_ptr()

298 + if region_addr in seen_region_ptrs:

299 + continue

300 + seen_region_ptrs.add(region_addr)

301 + base_addrs.append(region_addr)

302 + block_lens.append(cache.stride(0) * cache.element_size())

303 + self.kv_caches_base_addr = base_addrs

304 + self.block_len = block_lens

305

306 def prepare_values(

307 self,

308 chunks: Sequence[tuple[int, int]],

309 block_ids: list[int],

310 + shard_ids: Sequence[StoreShardId],

311 ) -> tuple[list[list[int]], list[list[int]], list[int]]:

216 - """Compute memory addresses and sizes for multiple token ranges.

217 -

218 - Returns:

219 - (addr_lists, size_lists, chunk_block_ids), one entry per chunk.

220 - """

312 if not chunks:

313 return [], [], []

314 + assert len(chunks) == len(shard_ids) and all(

315 + shard_id == 0 for shard_id in shard_ids

316 + )

317 base = np.asarray(self.kv_caches_base_addr, dtype=np.int64)

318 length = len(self.block_len)

319 blen = np.asarray(

@@ -229,7 +323,7 @@ class ChunkedTokenDatabase:

323 n = len(chunks)

324 starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n)

325 spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts

232 - assert not (spans % self.hash_block_size).any()

326 + assert np.all(spans % self.hash_block_size == 0)

327 bids = np.fromiter(

328 (block_ids[i] for i in (starts // self.block_size).tolist()),

329 dtype=np.int64,

@@ -241,15 +335,272 @@ class ChunkedTokenDatabase:

335 return addrs.tolist(), sizes.tolist(), bids.tolist()

336

337 def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]:

244 - """Return addresses and sizes for one physical block slot."""

245 - addr_list = []

246 - size_list = []

338 length = len(self.block_len)

248 - for index, base_addr in enumerate(self.kv_caches_base_addr):

249 - addr = base_addr + block_id * self.block_len[index % length]

250 - addr_list.append(addr)

251 - size_list.append(self.block_len[index % length])

252 - return addr_list, size_list

339 + return (

340 + [

341 + base_addr + block_id * self.block_len[index % length]

342 + for index, base_addr in enumerate(self.kv_caches_base_addr)

343 + ],

344 + [

345 + self.block_len[index % length]

346 + for index in range(len(self.kv_caches_base_addr))

347 + ],

348 + )

349 +

350 +

351 +class TPShardedStoreLayout(StoreLayout):

352 + """Store layout shared by divisible TP sizes."""

353 +

354 + store_format: str

355 +

356 + @classmethod

357 + def shared_namespace(cls, store_tp_size: int, pp_size: int) -> str:

358 + return (

359 + f"@store_tp:{store_tp_size}@store_pp:{pp_size}"

360 + f"@store_format:{cls.store_format}"

361 + )

362 +

363 + def __init__(

364 + self,

365 + metadata: KeyMetadata,

366 + block_size: int,

367 + hash_block_size: int,

368 + local_tp_size: int,

369 + store_tp_size: int,

370 + tp_rank: int,

371 + num_kv_heads: int,

372 + ) -> None:

373 + super().__init__(metadata, block_size, hash_block_size)

374 + self.shards_per_rank = store_tp_size // local_tp_size

375 + first_shard = tp_rank * self.shards_per_rank

376 + self.store_shard_ids = tuple(

377 + range(first_shard, first_shard + self.shards_per_rank)

378 + )

379 + self.store_tp_size = store_tp_size

380 + self.heads_per_store_shard = num_kv_heads // store_tp_size

381 + self.local_num_kv_heads = num_kv_heads // local_tp_size

382 + self._shard_key_prefixes = {

383 + shard_id: PoolKey.build_prefix(metadata, tp_rank=shard_id)

384 + for shard_id in self.store_shard_ids

385 + }

386 + self._shard_addr_bases = np.empty((self.shards_per_rank, 0), dtype=np.uint64)

387 + self._shard_block_strides = np.empty((self.shards_per_rank, 0), dtype=np.uint64)

388 + self._shard_sizes = np.empty((self.shards_per_rank, 0), dtype=np.uint64)

389 +

390 + @property

391 + def local_shard_ids(self) -> tuple[StoreShardId, ...]:

392 + return self.store_shard_ids

393 +

394 + def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:

395 + return PoolKey.build_key_string(

396 + self._shard_key_prefixes[shard_id], chunk_hash.hex()

397 + )

398 +

399 + def lookup_key_prefixes(

400 + self,

401 + rank_namespaces: Sequence[RankNamespace],

402 + ) -> tuple[str, ...]:

403 + pp_ranks = sorted({namespace[3] for namespace in rank_namespaces})

404 + return tuple(

405 + PoolKey.build_prefix(self.metadata, tp_rank=shard_id, pp_rank=pp_rank)

406 + for shard_id in range(self.store_tp_size)

407 + for pp_rank in pp_ranks

408 + )

409 +

410 + def _set_segment_templates(

411 + self, templates: Sequence[tuple[list[int], list[int], list[int]]]

412 + ) -> None:

413 + self._shard_addr_bases = np.asarray(

414 + [template[0] for template in templates], dtype=np.uint64

415 + )

416 + self._shard_block_strides = np.asarray(

417 + [template[1] for template in templates], dtype=np.uint64

418 + )

419 + self._shard_sizes = np.asarray(

420 + [template[2] for template in templates], dtype=np.uint64

421 + )

422 +

423 + def _has_packed_strides(

424 + self, head_stride: int, token_stride: int, content_bytes: int

425 + ) -> bool:

426 + raise NotImplementedError

427 +

428 + def _shard_segments(

429 + self, head_start: int, head_stride: int, token_stride: int, content_bytes: int

430 + ) -> Iterable[tuple[int, int]]:

431 + raise NotImplementedError

432 +

433 + def register_kv_caches(

434 + self,

435 + kv_caches: Sequence[torch.Tensor],

436 + num_blocks: int,

437 + ) -> None:

438 + templates: list[tuple[list[int], list[int], list[int]]] = [

439 + ([], [], []) for _ in range(self.shards_per_rank)

440 + ]

441 + for cache in kv_caches:

442 + if cache.ndim != 4 or tuple(cache.shape[:3]) != (

443 + num_blocks,

444 + self.local_num_kv_heads,

445 + self.block_size,

446 + ):

447 + raise ValueError(

448 + "TP-shared Mooncake store requires packed KV caches with "

449 + "logical shape (num_blocks, local_kv_heads, block_size, content)"

450 + )

451 +

452 + element_size = cache.element_size()

453 + block_stride, head_stride, token_stride, content_stride = (

454 + stride * element_size for stride in cache.stride()

455 + )

456 + content_bytes = cache.shape[3] * element_size

457 + if content_stride != element_size or not self._has_packed_strides(

458 + head_stride, token_stride, content_bytes

459 + ):

460 + raise ValueError(

461 + "TP-shared Mooncake store requires packed "

462 + f"{self.store_format} KV layout"

463 + )

464 +

465 + for local_shard, (addr_bases, block_strides, sizes) in enumerate(templates):

466 + head_start = local_shard * self.heads_per_store_shard

467 + for offset, size in self._shard_segments(

468 + head_start, head_stride, token_stride, content_bytes

469 + ):

470 + addr_bases.append(cache.data_ptr() + offset)

471 + block_strides.append(block_stride)

472 + sizes.append(size)

473 +

474 + self._set_segment_templates(templates)

475 +

476 + def prepare_values(

477 + self,

478 + chunks: Sequence[tuple[int, int]],

479 + block_ids: list[int],

480 + shard_ids: Sequence[StoreShardId],

481 + ) -> tuple[list[list[int]], list[list[int]], list[int]]:

482 + if not chunks:

483 + return [], [], []

484 + if len(chunks) != len(shard_ids):

485 + raise ValueError("Each Store chunk must have one shard ID")

486 +

487 + count = len(chunks)

488 + starts = np.fromiter(

489 + (chunk[0] for chunk in chunks), dtype=np.int64, count=count

490 + )

491 + ends = np.fromiter((chunk[1] for chunk in chunks), dtype=np.int64, count=count)

492 + if np.any(ends - starts != self.block_size):

493 + raise ValueError("TP-shared Mooncake store requires full KV blocks")

494 +

495 + first_shard = self.store_shard_ids[0]

496 + local_shards = np.fromiter(

497 + (shard_id - first_shard for shard_id in shard_ids),

498 + dtype=np.int64,

499 + count=count,

500 + )

501 + if np.any(local_shards < 0) or np.any(local_shards >= self.shards_per_rank):

502 + raise ValueError("Store shard is not owned by this TP rank")

503 + chunk_block_ids = np.fromiter(

504 + (block_ids[index] for index in (starts // self.block_size).tolist()),

505 + dtype=np.uint64,

506 + count=count,

507 + )

508 + addr_bases = self._shard_addr_bases[local_shards]

509 + block_strides = self._shard_block_strides[local_shards]

510 + addrs = addr_bases + chunk_block_ids[:, None] * block_strides

511 + sizes = self._shard_sizes[local_shards]

512 + return addrs.tolist(), sizes.tolist(), chunk_block_ids.tolist()

513 +

514 +

515 +class LBHNCStoreLayout(TPShardedStoreLayout):

516 + """Native head-major layout shared by divisible TP sizes."""

517 +

518 + store_format = "tp_shared_lbhnc"

519 +

520 + def _has_packed_strides(

521 + self, head_stride: int, token_stride: int, content_bytes: int

522 + ) -> bool:

523 + return (

524 + token_stride == content_bytes

525 + and head_stride == self.block_size * content_bytes

526 + )

527 +

528 + def _shard_segments(

529 + self, head_start: int, head_stride: int, token_stride: int, content_bytes: int

530 + ) -> Iterable[tuple[int, int]]:

531 + yield head_start * head_stride, self.heads_per_store_shard * head_stride

532 +

533 +

534 +class LBNHCStoreLayout(TPShardedStoreLayout):

535 + """Native token-major layout shared by divisible TP sizes."""

536 +

537 + store_format = "tp_shared_lbnhc"

538 +

539 + def _has_packed_strides(

540 + self, head_stride: int, token_stride: int, content_bytes: int

541 + ) -> bool:

542 + return (

543 + head_stride == content_bytes

544 + and token_stride == self.local_num_kv_heads * content_bytes

545 + )

546 +

547 + def _shard_segments(

548 + self, head_start: int, head_stride: int, token_stride: int, content_bytes: int

549 + ) -> Iterable[tuple[int, int]]:

550 + for token_idx in range(self.block_size):

551 + yield (

552 + token_idx * token_stride + head_start * head_stride,

553 + self.heads_per_store_shard * content_bytes,

554 + )

555 +

556 +

557 +class ChunkedTokenDatabase:

558 + """Enumerates logical token chunks and their hashes."""

559 +

560 + def __init__(

561 + self,

562 + metadata: KeyMetadata,

563 + block_size: int,

564 + hash_block_size: int | None = None,

565 + store_layout: StoreLayout | None = None,

566 + ):

567 + self.metadata = metadata

568 + self.block_size = block_size

569 + self.hash_block_size = hash_block_size or block_size

570 + if self.block_size % self.hash_block_size != 0:

571 + raise ValueError(

572 + f"block_size ({self.block_size}) must be a multiple of "

573 + f"hash_block_size ({self.hash_block_size})"

574 + )

575 + self.store_layout = store_layout or RankLocalStoreLayout(

576 + metadata, block_size, self.hash_block_size

577 + )

578 +

579 + @property

580 + def kv_caches_base_addr(self) -> list[int]:

581 + return self._rank_local_layout().kv_caches_base_addr

582 +

583 + @property

584 + def block_len(self) -> list[int]:

585 + return self._rank_local_layout().block_len

586 +

587 + def _rank_local_layout(self) -> RankLocalStoreLayout:

588 + if not isinstance(self.store_layout, RankLocalStoreLayout):

589 + raise RuntimeError("This operation requires a rank-local Store layout")

590 + return self.store_layout

591 +

592 + def key_for(self, chunk_hash: BlockHash) -> str:

593 + return self._rank_local_layout().key_for(0, chunk_hash)

594 +

595 + def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]):

596 + self._rank_local_layout().set_kv_caches_base_addr(kv_caches_base_addr)

597 +

598 + def set_block_len(self, block_len: list[int]):

599 + self._rank_local_layout().set_block_len(block_len)

600 +

601 + def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]:

602 + """Return addresses and sizes for one physical block slot."""

603 + return self._rank_local_layout().prepare_value_for_block(block_id)

604

605 def process_tokens(

606 self,

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+317 -176

@@ -12,6 +12,7 @@ and MooncakeDistributedStore integration.

13 import dataclasses

14 import json

15 +import math

16 import os

17 import queue

18 import socket

@@ -44,12 +45,16 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

45 BlobBlockHashes,

46 ChunkedTokenDatabase,

47 KeyMetadata,

48 + LBHNCStoreLayout,

49 + LBNHCStoreLayout,

50 MooncakeLookupResult,

51 MooncakeStoreConnectorMetadata,

52 MooncakeStoreWorkerMetadata,

53 PoolKey,

54 ReqMeta,

55 + StoreShardId,

56 TailKeyBoundary,

57 + TPShardedStoreLayout,

58 )

59 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import ( # noqa: E501

60 LOOKUP_MSG,

@@ -62,7 +67,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import

67 from vllm.logger import init_logger

68 from vllm.utils.math_utils import cdiv

69 from vllm.utils.network_utils import get_ip, make_zmq_socket

65 -from vllm.utils.torch_utils import is_non_overlapping_and_dense

70 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

71 from vllm.v1.core.kv_cache_utils import (

72 BlockHash,

@@ -71,6 +75,7 @@ from vllm.v1.core.kv_cache_utils import (

75 resolve_kv_cache_block_sizes,

76 )

77 from vllm.v1.kv_cache_interface import (

78 + FullAttentionSpec,

79 KVCacheConfig,

80 KVCacheSpec,

81 MambaSpec,

@@ -79,6 +84,7 @@ from vllm.v1.kv_cache_interface import (

84 UniformTypeKVCacheSpecs,

85 group_kernel_blocks,

86 )

87 +from vllm.v1.kv_cache_layout import KVCacheLayout

89 from .metrics import MooncakeStoreConnectorStats

@@ -92,6 +98,22 @@ MOONCAKE_NO_AVAILABLE_HANDLE = -200

98 _T = TypeVar("_T")

100

101 +def resolve_store_tp_size(extra_config: dict[str, Any]) -> int | None:

102 + """Resolve the common Store TP requested by connector config."""

103 + if extra_config.get("enable_store_tp_lcm") is True:

104 + prefill_tp_sizes = extra_config.get("prefill_tp_sizes")

105 + if not isinstance(prefill_tp_sizes, list) or not prefill_tp_sizes:

106 + return None

107 + if any(

108 + type(tp_size) is not int or tp_size <= 0 for tp_size in prefill_tp_sizes

109 + ):

110 + return None

111 + return math.lcm(*prefill_tp_sizes)

112 +

113 + store_tp_size = extra_config.get("store_tp_size")

114 + return store_tp_size if type(store_tp_size) is int and store_tp_size > 0 else None

115 +

116 +

117 def _rotate_list(values: list[_T], offset: int) -> list[_T]:

118 return values[offset:] + values[:offset]

119

@@ -109,7 +131,10 @@ def _make_mooncake_group_id(metadata: KeyMetadata, chunk_hash: str) -> str:

131 # Mooncake group ids describe the lifecycle unit. For vLLM, that unit is

132 # a prefix chunk, so shard dimensions stay only in the object key.

133 prefix = f"{metadata.cache_prefix}@" if metadata.cache_prefix else ""

112 - return f"vllm-mooncake-store:{prefix}{metadata.model_name}@{chunk_hash}"

134 + return (

135 + f"vllm-mooncake-store:{prefix}{metadata.model_name}"

136 + f"{metadata.store_namespace}@{chunk_hash}"

137 + )

138

139

140 # Mirrors FileStorageConfig::local_buffer_size in Mooncake C++.

@@ -952,8 +977,11 @@ class KVCacheStoreSendingThread(KVTransferThread):

977 starts: list[int] = []

978 ends: list[int] = []

979 keys: list[str] = []

955 - kv_event_block_hashes: list[BlockHash] = []

980 + event_specs: list[tuple[int, int, int, BlockHash]] | None = (

981 + [] if self.enable_kv_event else None

982 + )

983 group_indices: list[int] = []

984 + store_shard_ids: list[StoreShardId] = []

985 for g_idx, db in enumerate(self.token_databases):

986 # Rotate the stride phase per group to balance load across ranks.

987 put_step = self.group_put_steps[g_idx]

@@ -980,12 +1008,14 @@ class KVCacheStoreSendingThread(KVTransferThread):

1008 block_idx,

1009 )

1010 continue

983 - starts.append(start)

984 - ends.append(end)

985 - keys.append(db.key_for(block_hash))

986 - if self.enable_kv_event:

987 - kv_event_block_hashes.append(block_hash)

988 - group_indices.append(g_idx)

1011 + for store_shard_id in db.store_layout.local_shard_ids:

1012 + starts.append(start)

1013 + ends.append(end)

1014 + keys.append(db.store_layout.key_for(store_shard_id, block_hash))

1015 + group_indices.append(g_idx)

1016 + store_shard_ids.append(store_shard_id)

1017 + if event_specs is not None:

1018 + event_specs.append((start, end, g_idx, block_hash))

1019

1020 if not keys:

1021 self._record_saved(req_meta, token_len)

@@ -1023,11 +1053,10 @@ class KVCacheStoreSendingThread(KVTransferThread):

1053 starts = [starts[i] for i in missing_indices]

1054 ends = [ends[i] for i in missing_indices]

1055 keys = [keys[i] for i in missing_indices]

1026 - if self.enable_kv_event:

1027 - kv_event_block_hashes = [

1028 - kv_event_block_hashes[i] for i in missing_indices

1029 - ]

1056 + if event_specs is not None:

1057 + event_specs = [event_specs[i] for i in missing_indices]

1058 group_indices = [group_indices[i] for i in missing_indices]

1059 + store_shard_ids = [store_shard_ids[i] for i in missing_indices]

1060

1061 group_ids = (

1062 [

@@ -1050,61 +1079,33 @@ class KVCacheStoreSendingThread(KVTransferThread):

1079

1080 addrs: list[list[int]] = []

1081 sizes: list[list[int]] = []

1053 - stored_events: list[BlockStored] = []

1082 chunks_per_group: list[list[tuple[int, int]]] = [

1083 [] for _ in self.token_databases

1084 ]

1057 - for start, end, g_idx in zip(starts, ends, group_indices, strict=True):

1085 + shards_per_group: list[list[StoreShardId]] = [

1086 + [] for _ in self.token_databases

1087 + ]

1088 + for start, end, g_idx, store_shard_id in zip(

1089 + starts,

1090 + ends,

1091 + group_indices,

1092 + store_shard_ids,

1093 + strict=True,

1094 + ):

1095 chunks_per_group[g_idx].append((start, end))

1096 + shards_per_group[g_idx].append(store_shard_id)

1097 for g_idx, chunks in enumerate(chunks_per_group):

1098 if not chunks:

1099 continue

1100 db = self.token_databases[g_idx]

1063 - group_addrs, group_sizes, _ = db.prepare_values(

1064 - chunks, block_ids_per_group[g_idx]

1101 + group_addrs, group_sizes, _ = db.store_layout.prepare_values(

1102 + chunks,

1103 + block_ids_per_group[g_idx],

1104 + shards_per_group[g_idx],

1105 )

1106 addrs.extend(group_addrs)

1107 sizes.extend(group_sizes)

1108

1069 - if self.enable_kv_event:

1070 - new_block_hashes = [

1071 - maybe_convert_block_hash(bh) for bh in kv_event_block_hashes

1072 - ]

1073 - token_ids_end = token_ids_start + len(event_token_ids or ())

1074 -

1075 - for idx, (s, e, g_idx) in enumerate(

1076 - zip(starts, ends, group_indices, strict=True)

1077 - ):

1078 - db = self.token_databases[g_idx]

1079 - if self.enable_kv_event:

1080 - token_ids = (

1081 - event_token_ids[s - token_ids_start : e - token_ids_start]

1082 - if event_token_ids is not None

1083 - and token_ids_start <= s

1084 - and e <= token_ids_end

1085 - else []

1086 - )

1087 - stored_event = BlockStored(

1088 - block_hashes=[new_block_hashes[idx]],

1089 - # Derive the direct predecessor from the unfiltered

1090 - # request chain. Adjacent PUTs need not be adjacent in

1091 - # that chain after Store dedup, masks, or TP striding.

1092 - parent_block_hash=(

1093 - maybe_convert_block_hash(

1094 - req_meta.block_hashes[s // db.hash_block_size - 1]

1095 - )

1096 - if s > 0

1097 - else None

1098 - ),

1099 - token_ids=token_ids,

1100 - block_size=db.block_size,

1101 - lora_id=None,

1102 - medium="cpu",

1103 - lora_name=None,

1104 - group_idx=g_idx,

1105 - )

1106 - stored_events.append(stored_event)

1107 -

1109 if current_event is not None:

1110 current_event.synchronize()

1111

@@ -1112,6 +1113,8 @@ class KVCacheStoreSendingThread(KVTransferThread):

1113 assert len(group_ids) == len(keys)

1114 self.replicate_config.group_ids = group_ids

1115

1116 + failed_indices: set[int] = set()

1117 + put_had_exception = False

1118 batch_bytes = _sum_batch_bytes(sizes)

1119 put_start = time.perf_counter()

1120 try:

@@ -1121,65 +1124,106 @@ class KVCacheStoreSendingThread(KVTransferThread):

1124 sizes,

1125 self.replicate_config,

1126 )

1124 - failed = [i for i, v in enumerate(res) if v < 0]

1127 + except Exception as e:

1128 self._record_operation(

1129 "save_put",

1130 put_start,

1131 len(keys),

1132 num_bytes=batch_bytes,

1130 - status="partial_failure" if failed else "ok",

1131 - num_failed_keys=len(failed),

1133 + status="error",

1134 + num_failed_keys=len(keys),

1135 )

1133 - if failed:

1134 - failed_codes = set(res[i] for i in failed)

1135 - if self.enable_kv_event:

1136 - failed_indices = set(failed)

1137 - stored_events = [

1138 - event

1139 - for i, event in enumerate(stored_events)

1140 - if i not in failed_indices

1141 - ]

1136 + logger.error("Failed to put key %s, error: %s", keys, e)

1137 + put_had_exception = True

1138 + else:

1139 + failed_indices = {i for i, value in enumerate(res) if value < 0}

1140 + self._record_operation(

1141 + "save_put",

1142 + put_start,

1143 + len(keys),

1144 + num_bytes=batch_bytes,

1145 + status="partial_failure" if failed_indices else "ok",

1146 + num_failed_keys=len(failed_indices),

1147 + )

1148 + failed_codes = {res[i] for i in failed_indices}

1149 + if failed_indices:

1150 logger.warning(

1151 "batch_put failed: %d/%d keys failed "

1144 - "(codes=%s, batch_bytes=%d, num_keys=%d), "

1145 - "first_key=%s",

1146 - len(failed),

1152 + "(codes=%s, batch_bytes=%d), first_key=%s",

1153 + len(failed_indices),

1154 len(keys),

1155 failed_codes,

1156 batch_bytes,

1150 - len(keys),

1151 - keys[0] if keys else "N/A",

1157 + keys[0],

1158 )

1153 - if (

1154 - MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes

1155 - and not self._mark_request_skipped_for_pressure(req_meta)

1156 - ):

1157 - logger.warning(

1158 - "Detected Mooncake CPU/disk offloading pressure "

1159 - "(NO_AVAILABLE_HANDLE); skipping future store "

1160 - "batches for request %s until a later store "

1161 - "batch succeeds",

1162 - req_id,

1163 - )

1164 - else:

1165 - self._record_saved(req_meta, token_len)

1166 - save_completed = True

1167 - if self._clear_store_pressure():

1168 - logger.info(

1169 - "Mooncake CPU/disk offloading pressure cleared "

1170 - "after a successful store batch"

1159 + if (

1160 + MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes

1161 + and not self._mark_request_skipped_for_pressure(req_meta)

1162 + ):

1163 + logger.warning(

1164 + "Detected Mooncake CPU/disk offloading pressure "

1165 + "(NO_AVAILABLE_HANDLE); skipping future store "

1166 + "batches for request %s until a later store batch succeeds",

1167 + req_id,

1168 + )

1169 +

1170 + if not put_had_exception and not failed_indices:

1171 + self._record_saved(req_meta, token_len)

1172 + save_completed = True

1173 + if self._clear_store_pressure():

1174 + logger.info(

1175 + "Mooncake CPU/disk offloading pressure cleared "

1176 + "after a successful store batch"

1177 + )

1178 +

1179 + stored_events: list[BlockStored] = []

1180 + if self.enable_kv_event and not put_had_exception:

1181 + assert event_specs is not None

1182 + # BlockStored is a logical-block event, while one block may map

1183 + # to several Store shards. Emit once only after every missing

1184 + # shard for that block succeeded. Shards that already existed

1185 + # were removed before this mapping and are already satisfied.

1186 + indices_by_event: dict[tuple[int, int, int, BlockHash], list[int]] = {}

1187 + for index, event_spec in enumerate(event_specs):

1188 + indices_by_event.setdefault(event_spec, []).append(index)

1189 +

1190 + token_ids_end = token_ids_start + len(event_token_ids or ())

1191 + for (

1192 + s,

1193 + end,

1194 + g_idx,

1195 + block_hash,

1196 + ), event_indices in indices_by_event.items():

1197 + if any(index in failed_indices for index in event_indices):

1198 + continue

1199 + db = self.token_databases[g_idx]

1200 + token_ids = (

1201 + event_token_ids[s - token_ids_start : end - token_ids_start]

1202 + if event_token_ids is not None

1203 + and token_ids_start <= s

1204 + and end <= token_ids_end

1205 + else []

1206 + )

1207 + stored_events.append(

1208 + BlockStored(

1209 + block_hashes=[maybe_convert_block_hash(block_hash)],

1210 + # Store filtering can separate adjacent request

1211 + # blocks, so derive the predecessor from the request.

1212 + parent_block_hash=(

1213 + maybe_convert_block_hash(

1214 + req_meta.block_hashes[s // db.hash_block_size - 1]

1215 + )

1216 + if s > 0

1217 + else None

1218 + ),

1219 + token_ids=token_ids,

1220 + block_size=db.block_size,

1221 + lora_id=None,

1222 + medium="cpu",

1223 + lora_name=None,

1224 + group_idx=g_idx,

1225 )

1172 - except Exception as e:

1173 - self._record_operation(

1174 - "save_put",

1175 - put_start,

1176 - len(keys),

1177 - num_bytes=batch_bytes,

1178 - status="error",

1179 - num_failed_keys=len(keys),

1180 - )

1181 - logger.error("Failed to put key %s, error: %s", keys, e)

1182 - stored_events.clear()

1226 + )

1227

1228 if self.enable_kv_event and stored_events:

1229 self.update_kv_event(stored_events)

@@ -1269,6 +1313,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1313 for g_idx, db in enumerate(self.token_databases):

1314 mask = load_mask_per_group[g_idx]

1315 chunks: list[tuple[int, int]] = []

1316 + store_shard_ids: list[StoreShardId] = []

1317 for start, end, block_hash in db.process_tokens(

1318 token_len, req_meta.block_hashes, mask_num

1319 ):

@@ -1282,10 +1327,14 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1327 block_hash = req_meta.block_hashes[

1328 boundary_tokens // db.hash_block_size - 1

1329 ]

1285 - key_list.append(db.key_for(block_hash))

1286 - chunks.append((start, end))

1287 - g_addrs, g_sizes, g_block_ids = db.prepare_values(

1288 - chunks, req_meta.block_ids[g_idx]

1330 + for store_shard_id in db.store_layout.local_shard_ids:

1331 + key_list.append(db.store_layout.key_for(store_shard_id, block_hash))

1332 + chunks.append((start, end))

1333 + store_shard_ids.append(store_shard_id)

1334 + g_addrs, g_sizes, g_block_ids = db.store_layout.prepare_values(

1335 + chunks,

1336 + req_meta.block_ids[g_idx],

1337 + store_shard_ids,

1338 )

1339 addr_list.extend(g_addrs)

1340 size_list.extend(g_sizes)

@@ -1298,7 +1347,14 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1347 size_list_c = _rotate_list(size_list, rotation)

1348 block_id_list_c = _rotate_list(block_id_list, rotation)

1349

1301 - load_batches = [(key_list_c, addr_list_c, size_list_c, block_id_list_c)]

1350 + load_batches = [

1351 + (

1352 + key_list_c,

1353 + addr_list_c,

1354 + size_list_c,

1355 + block_id_list_c,

1356 + )

1357 + ]

1358 if self.usable_disk_offload_buffer_budget_bytes is not None:

1359 total_staging_bytes = sum(

1360 _estimate_disk_offload_staging_bytes(size) for size in size_list_c

@@ -1340,7 +1396,12 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1396 block_id_offset:next_block_id_offset

1397 ]

1398 load_batches.append(

1343 - (batch_keys, batch_addrs, batch_sizes, batch_block_ids)

1399 + (

1400 + batch_keys,

1401 + batch_addrs,

1402 + batch_sizes,

1403 + batch_block_ids,

1404 + )

1405 )

1406 block_id_offset = next_block_id_offset

1407

@@ -1348,7 +1409,12 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1409 current_batch_block_ids: list[int] = block_id_list_c

1410 batch_bytes = 0

1411 try:

1351 - for batch_keys, batch_addrs, batch_sizes, batch_block_ids in load_batches:

1412 + for (

1413 + batch_keys,

1414 + batch_addrs,

1415 + batch_sizes,

1416 + batch_block_ids,

1417 + ) in load_batches:

1418 current_batch_keys = batch_keys

1419 current_batch_block_ids = batch_block_ids

1420 batch_bytes = _sum_batch_bytes(batch_sizes)

@@ -1607,11 +1673,9 @@ class MooncakeStoreWorker:

1673 retention_interval=kv_cache_config.prefix_cache_retention_interval,

1674 dcp_world_size=self.dcp_size,

1675 )

1610 - # One ChunkedTokenDatabase per group; addresses populated in

1611 - # register_kv_caches once the kv-cache layout is known. Each group's

1612 - # key namespace is its TP shard id: ranks holding identical bytes

1613 - # (MLA / shared GQA KV heads) share a namespace, TP-sharded Mamba

1614 - # state gets one namespace per rank.

1676 + self.store_tp_size, store_namespace, store_layout_cls = (

1677 + self._select_store_layout(extra_config)

1678 + )

1679 metadata = KeyMetadata(

1680 model_name=model_config.model.rstrip("/").split("/")[-1],

1681 tp_rank=self.tp_rank,

@@ -1623,24 +1687,140 @@ class MooncakeStoreWorker:

1687 "cache_prefix", ""

1688 )

1689 ),

1690 + store_namespace=store_namespace,

1691 )

1692 self._group_tp_replication_factors: tuple[int, ...] = (

1693 self._compute_group_tp_replication_factors()

1694 )

1630 - self.token_dbs = [

1631 - ChunkedTokenDatabase(

1632 - dataclasses.replace(

1633 - metadata,

1634 - group_id=g_idx,

1635 - tp_rank=self.tp_rank // self._group_tp_replication_factors[g_idx],

1636 - ),

1637 - g.kv_cache_spec.block_size,

1638 - hash_block_size=self.hash_block_size,

1639 - )

1640 - for g_idx, g in enumerate(self._kv_cache_groups)

1641 - ]

1695 + self.token_dbs = self._build_token_databases(metadata, store_layout_cls)

1696 self._init_lookup_key_prefixes()

1697

1698 + def _supports_tp_sharded_store_layout(

1699 + self,

1700 + layout_cls: type[TPShardedStoreLayout] | None,

1701 + extra_config: dict[str, Any],

1702 + ) -> bool:

1703 + if (

1704 + layout_cls is None

1705 + or self.pcp_size != 1

1706 + or self.dcp_size != 1

1707 + or len(self._kv_cache_groups) != 1

1708 + ):

1709 + return False

1710 +

1711 + # Subclasses may use different sharding or cache-lifetime semantics.

1712 + return (

1713 + type(self._kv_cache_groups[0].kv_cache_spec) is FullAttentionSpec

1714 + and str(extra_config.get("enable_cross_layers_blocks", "False")).lower()

1715 + != "true"

1716 + )

1717 +

1718 + def _select_store_layout(

1719 + self, extra_config: dict[str, Any]

1720 + ) -> tuple[int | None, str, type[TPShardedStoreLayout] | None]:

1721 + """Select the opt-in TP layout and its Store namespace."""

1722 + lcm_store_tp_enabled = extra_config.get("enable_store_tp_lcm") is True

1723 + store_tp_requested = (

1724 + lcm_store_tp_enabled or extra_config.get("store_tp_size") is not None

1725 + )

1726 + if not store_tp_requested:

1727 + return None, "", None

1728 +

1729 + requested_store_tp_size = resolve_store_tp_size(extra_config)

1730 + cache_layout = self.cache_config.get_resolved_kv_cache_layout()

1731 + layout_cls: type[TPShardedStoreLayout] | None = {

1732 + KVCacheLayout.LBHNC: LBHNCStoreLayout,

1733 + KVCacheLayout.LBNHC: LBNHCStoreLayout,

1734 + }.get(cache_layout)

1735 +

1736 + if (

1737 + requested_store_tp_size is not None

1738 + and requested_store_tp_size >= self.tp_size

1739 + and requested_store_tp_size % self.tp_size == 0

1740 + and self._supports_tp_sharded_store_layout(layout_cls, extra_config)

1741 + ):

1742 + assert layout_cls is not None

1743 + if self.num_kv_head % requested_store_tp_size == 0:

1744 + if layout_cls is LBNHCStoreLayout:

1745 + logger.warning_once(

1746 + "Mooncake Store TP sharding is using the LBNHC (NHC) "

1747 + "KV cache layout, which creates many transfer segments "

1748 + "and may significantly reduce PUT/GET performance. Use "

1749 + "LBHNC (HNC) when supported by the attention backend."

1750 + )

1751 + return (

1752 + requested_store_tp_size,

1753 + layout_cls.shared_namespace(requested_store_tp_size, self.pp_size),

1754 + layout_cls,

1755 + )

1756 + if self.num_kv_head == 1:

1757 + logger.info(

1758 + "Mooncake heterogeneous-TP store sharing uses the replicated "

1759 + "MQA layout for store_tp_size=%d",

1760 + requested_store_tp_size,

1761 + )

1762 + return (

1763 + None,

1764 + f"@store_pp:{self.pp_size}@store_format:tp_shared_mqa",

1765 + None,

1766 + )

1767 +

1768 + store_namespace = (

1769 + f"@store_pp:{self.pp_size}@store_format:"

1770 + f"rank_local_tp{self.tp_size}_layout_{cache_layout.name}"

1771 + )

1772 + requested_topology = (

1773 + extra_config.get("prefill_tp_sizes")

1774 + if lcm_store_tp_enabled

1775 + else extra_config.get("store_tp_size")

1776 + )

1777 + logger.warning(

1778 + "Mooncake heterogeneous-TP store sharing is disabled for "

1779 + "Store TP configuration %r with KV layout %s; using a "

1780 + "compatibility-namespaced rank-local store layout",

1781 + requested_topology,

1782 + cache_layout,

1783 + )

1784 + return None, store_namespace, None

1785 +

1786 + def _build_token_databases(

1787 + self,

1788 + metadata: KeyMetadata,

1789 + layout_cls: type[TPShardedStoreLayout] | None,

1790 + ) -> list[ChunkedTokenDatabase]:

1791 + """Construct token databases and their Store layouts."""

1792 + token_dbs: list[ChunkedTokenDatabase] = []

1793 + for group_idx, group in enumerate(self._kv_cache_groups):

1794 + group_tp_rank = self.tp_rank

1795 + if layout_cls is None:

1796 + group_tp_rank //= self._group_tp_replication_factors[group_idx]

1797 + group_metadata = dataclasses.replace(

1798 + metadata,

1799 + group_id=group_idx,

1800 + tp_rank=group_tp_rank,

1801 + )

1802 + store_layout: TPShardedStoreLayout | None = None

1803 + if layout_cls is not None:

1804 + assert self.store_tp_size is not None

1805 + store_layout = layout_cls(

1806 + group_metadata,

1807 + group.kv_cache_spec.block_size,

1808 + self.hash_block_size,

1809 + local_tp_size=self.tp_size,

1810 + store_tp_size=self.store_tp_size,

1811 + tp_rank=self.tp_rank,

1812 + num_kv_heads=self.num_kv_head,

1813 + )

1814 + token_dbs.append(

1815 + ChunkedTokenDatabase(

1816 + group_metadata,

1817 + group.kv_cache_spec.block_size,

1818 + hash_block_size=self.hash_block_size,

1819 + store_layout=store_layout,

1820 + )

1821 + )

1822 + return token_dbs

1823 +

1824 def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int:

1825 if self.dcp_size > 1:

1826 return 1

@@ -1689,17 +1869,8 @@ class MooncakeStoreWorker:

1869 )

1870

1871 self._lookup_key_prefixes = tuple(

1692 - tuple(

1693 - PoolKey.build_prefix(

1694 - db.metadata,

1695 - tp_rank=tp_rank,

1696 - pcp_rank=pcp_rank,

1697 - dcp_rank=dcp_rank,

1698 - pp_rank=pp_rank,

1699 - )

1700 - for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces(

1701 - self._group_tp_replication_factors[g_idx]

1702 - )

1872 + db.store_layout.lookup_key_prefixes(

1873 + rank_namespaces(self._group_tp_replication_factors[g_idx])

1874 )

1875 for g_idx, db in enumerate(self.token_dbs)

1876 )

@@ -1719,12 +1890,11 @@ class MooncakeStoreWorker:

1890 self.num_blocks = self.cache_config.num_gpu_blocks

1891

1892 seen_storage_ptrs: set[int] = set()

1722 - seen_region_ptrs: set[int] = set()

1723 - addrs: list[int] = []

1724 - block_lens: list[int] = []

1893 + cache_tensors: list[torch.Tensor] = []

1894

1895 for cache in kv_caches.values():

1896 cache = group_kernel_blocks(cache, self.num_blocks)

1897 + cache_tensors.append(cache)

1898 cache_storage = cache.untyped_storage()

1899 base_addr = cache_storage.data_ptr()

1900 region_len = cache_storage.nbytes()

@@ -1740,48 +1910,15 @@ class MooncakeStoreWorker:

1910 ret,

1911 )

1912

1743 - if not is_non_overlapping_and_dense(cache[0]):

1744 - # A block is scattered across per-head regions; each region's

1745 - # blocks are contiguous.

1746 - for head_idx in range(cache.shape[1]):

1747 - head_cache = cache[:, head_idx]

1748 - assert is_non_overlapping_and_dense(head_cache[0])

1749 - region_addr = head_cache.data_ptr()

1750 - if region_addr in seen_region_ptrs:

1751 - continue

1752 - seen_region_ptrs.add(region_addr)

1753 - addrs.append(region_addr)

1754 - block_lens.append(head_cache.stride(0) * head_cache.element_size())

1755 - elif cache.stride(0) * cache.element_size() * self.num_blocks == region_len:

1756 - # The block stride spans the whole per-block window

1757 - # (block-outermost and packed layouts, and single-layer

1758 - # tensors), which may hold other layers' pages at higher

1759 - # offsets. Register the storage once as one whole-window

1760 - # region; per-layer regions would copy overlapping windows and

1761 - # run past the storage for offset layers.

1762 - if base_addr in seen_region_ptrs:

1763 - continue

1764 - seen_region_ptrs.add(base_addr)

1765 - addrs.append(base_addr)

1766 - block_lens.append(region_len // self.num_blocks)

1767 - else:

1768 - region_addr = cache.data_ptr()

1769 - if region_addr in seen_region_ptrs:

1770 - continue

1771 - seen_region_ptrs.add(region_addr)

1772 - addrs.append(region_addr)

1773 - block_lens.append(cache.stride(0) * cache.element_size())

1774 -

1913 logger.info(

1776 - "Registered KV caches: num_groups=%d, num_segments=%d, num_blocks=%d",

1914 + "Registered KV caches: num_groups=%d, num_tensors=%d, num_blocks=%d",

1915 len(self.token_dbs),

1778 - len(addrs),

1916 + len(cache_tensors),

1917 self.num_blocks,

1918 )

1919

1920 for db in self.token_dbs:

1783 - db.set_kv_caches_base_addr(addrs)

1784 - db.set_block_len(block_lens)

1921 + db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)

1922

1923 # Start transfer threads

1924 if self.can_put:

@@ -2128,6 +2265,10 @@ class MooncakeStoreWorker:

2265 return self.kv_send_thread.get_kv_events()

2266 return []

2267

2268 + @property

2269 + def group_tp_replication_factors(self) -> tuple[int, ...]:

2270 + return self._group_tp_replication_factors

2271 +

2272 def close(self) -> None:

2273 """Release the MooncakeDistributedStore handle on teardown.

2274