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

[1/N][KV Connector] Identify externally transferable KV cache groups (#53779)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Aug 27, 2026 at 11:21 UTC d85708f7a4362334eb58cff1fdf5265b49cac310

11 files changed +137 -29

tests/v1/core/test_kv_cache_utils.py

+24

@@ -140,6 +140,30 @@ def new_kv_cache_spec(

140 )

141

142

143 +def test_kv_cache_config_selects_only_transferable_groups():

144 + """Connectors must see a stable projection of transfer-eligible groups."""

145 + groups = [

146 + KVCacheGroupSpec(["layer.0"], new_kv_cache_spec()),

147 + KVCacheGroupSpec(["layer.1"], new_kv_cache_spec(), enable_kv_transfer=False),

148 + KVCacheGroupSpec(["layer.2"], new_kv_cache_spec()),

149 + ]

150 + config = KVCacheConfig(

151 + num_blocks=1,

152 + kv_cache_tensors=[],

153 + kv_cache_groups=groups,

154 + )

155 +

156 + assert config.transfer_group_ids == (0, 2)

157 + assert config.transfer_groups == (groups[0], groups[2])

158 + assert config.transfer_group_index_by_layer == {"layer.0": 0, "layer.2": 1}

159 + first_blocks = [1, 2]

160 + third_blocks = [4]

161 + assert config.select_transfer_block_ids((first_blocks, [3], third_blocks)) == (

162 + first_blocks,

163 + third_blocks,

164 + )

165 +

166 +

167 def new_sliding_window_spec(

168 block_size=16,

169 num_kv_heads=2,

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

+17

@@ -33,6 +33,9 @@ def _make_bare_scheduler(

33 scheduler._block_size = 16

34 scheduler._hash_block_size = hash_block_size

35 scheduler.enable_partial_hash_hits = enable_partial_hash_hits

36 + scheduler.kv_cache_config = SimpleNamespace(

37 + select_transfer_block_ids=lambda block_ids: tuple(block_ids)

38 + )

39 scheduler.load_specs = {}

40 scheduler._unfinished_request_ids = {"req-0"}

41 scheduler._unfinished_requests = {}

@@ -179,6 +182,20 @@ def _add_unfinished_request(

182 )

183

184

185 +def test_update_state_excludes_nontransfer_groups():

186 + """Store metadata must match the worker's registered cache groups."""

187 + scheduler = _make_bare_scheduler()

188 + scheduler.kv_cache_config = SimpleNamespace(

189 + select_transfer_block_ids=lambda block_ids: (block_ids[0],)

190 + )

191 + request = SimpleNamespace(request_id="req-1")

192 + blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9]))

193 +

194 + scheduler.update_state_after_alloc(request, blocks, num_external_tokens=32)

195 +

196 + assert scheduler._unfinished_requests["req-1"][1] == ([1, 2],)

197 +

198 +

199 def _setup_decode_request(

200 *,

201 kv_role: str = "kv_consumer",

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

+14 -1

@@ -854,7 +854,10 @@ def test_post_process_zeroes_untransferred_tail():

854 worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache}

855 fa_group = MagicMock(layer_names=["attn.0"])

856 ssm_group = MagicMock(layer_names=["mamba.0"])

857 - worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group])

857 + worker.kv_cache_config = MagicMock(

858 + kv_cache_groups=[fa_group, ssm_group],

859 + transfer_groups=[fa_group, ssm_group],

860 + )

861 # The cached property filters mamba layers out of the permuted caches.

862 attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker)

863 assert len(attn_caches) == 1 and attn_caches[0] is attn_cache

@@ -1477,6 +1480,16 @@ def test_exchange_clipped_blocks_ssm_positional_states():

1480 assert clipped == ([1, 2, 3], [0, 5, 6, 7])

1481

1482

1483 +@pytest.mark.cpu_test

1484 +def test_exchange_clipped_blocks_excludes_nontransfer_groups():

1485 + """Scheduler block tables must match the worker's registered regions."""

1486 + sched = make_nixl_scheduler()

1487 + sched.kv_cache_config = make_kv_cache_config(block_size=16, swa_enabled=True)

1488 + sched.kv_cache_config.kv_cache_groups[1].enable_kv_transfer = False

1489 +

1490 + assert sched.get_exchange_clipped_blocks(([1, 2], [9])) == ([1, 2],)

1491 +

1492 +

1493 # ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────

1494

1495

tests/v1/kv_connector/unit/utils.py

@@ -524,6 +524,10 @@ def make_nixl_scheduler(

524 sched = object.__new__(NixlConnectorScheduler)

525 sched._has_mamba = has_mamba

526 sched._is_hma_required = is_hma_required

527 + sched.kv_cache_config = make_kv_cache_config(

528 + block_size=16,

529 + mamba_enabled=has_mamba,

530 + )

531

532 if heartbeat:

533 sched._heartbeat_by_engine = {}

@@ -582,6 +586,10 @@ def make_nixl_push_scheduler(

586 sched.side_channel_port = 5600

587 sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled

588 sched._has_mamba = has_mamba

589 + sched.kv_cache_config = make_kv_cache_config(

590 + block_size=16,

591 + mamba_enabled=has_mamba,

592 + )

593

594 # vllm_config is consulted for parallel_config.tensor_parallel_size.

595 vllm_config = MagicMock()

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

+2 -2

@@ -97,7 +97,7 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

98 unsupported: list[str] = []

99 cache_block_size = vllm_config.cache_config.block_size

100 - for g_idx, g in enumerate(kv_cache_config.kv_cache_groups):

100 + for g_idx, g in enumerate(kv_cache_config.transfer_groups):

101 spec = g.kv_cache_spec

102 if isinstance(spec, CrossAttentionSpec):

103 unsupported.append(f"group {g_idx}: CrossAttentionSpec")

@@ -110,7 +110,7 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

110 )

111 pcp = vllm_config.parallel_config.prefill_context_parallel_size

112 dcp = vllm_config.parallel_config.decode_context_parallel_size

113 - if len(kv_cache_config.kv_cache_groups) > 1 and pcp * dcp > 1:

113 + if len(kv_cache_config.transfer_groups) > 1 and pcp * dcp > 1:

114 unsupported.append(

115 f"PCP/DCP > 1 (pcp={pcp}, dcp={dcp}) with hybrid attention"

116 )

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

+15 -11

@@ -69,6 +69,7 @@ class MooncakeStoreScheduler:

69 kv_event_config and kv_event_config.enable_kv_cache_events

70 )

71 self.client = LookupKeyClient(vllm_config)

72 + self.kv_cache_config = kv_cache_config

74 # Align with the engine's own scheduler_block_size and hash_block_size.

75 self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(

@@ -157,7 +158,9 @@ class MooncakeStoreScheduler:

158 """Update state after block allocation."""

159 local_block_ids: tuple[list[int], ...] = ()

160 if num_external_tokens > 0:

160 - local_block_ids = blocks.get_block_ids()

161 + local_block_ids = self.kv_cache_config.select_transfer_block_ids(

162 + blocks.get_block_ids()

163 + )

164

165 self._unfinished_requests[request.request_id] = (request, local_block_ids)

166 self._unfinished_request_ids.add(request.request_id)

@@ -220,12 +223,12 @@ class MooncakeStoreScheduler:

223 request_tuple = self._unfinished_requests.get(request.req_id)

224 request_real = request_tuple[0] # type: ignore[index]

225

223 - if isinstance(request.block_ids, tuple):

224 - # Multi-group: preserve per-group structure.

225 - unfolded_block_ids = tuple(b.copy() for b in request.block_ids)

226 - else:

227 - # Single-group legacy: list[int] -> 1-tuple.

228 - unfolded_block_ids = (request.block_ids.copy(),)

226 + unfolded_block_ids = tuple(

227 + blocks.copy()

228 + for blocks in self.kv_cache_config.select_transfer_block_ids(

229 + request.block_ids

230 + )

231 + )

232

233 prefill_tokens = _new_req_prefill_tokens(request)

234 request_tracker = RequestTracker(

@@ -259,16 +262,17 @@ class MooncakeStoreScheduler:

262 if can_process_cached:

263 for i, req_id in enumerate(cached_reqs.req_ids):

264 new_block_ids = cached_reqs.new_block_ids[i]

265 + if new_block_ids:

266 + new_block_ids = self.kv_cache_config.select_transfer_block_ids(

267 + new_block_ids

268 + )

269

270 req_meta = None

271 if req_id in cached_reqs.resumed_req_ids:

272 # Resumed after preemption

273 if not new_block_ids:

274 continue

268 - if isinstance(new_block_ids, tuple):

269 - new_block_ids = tuple(b.copy() for b in new_block_ids)

270 - else:

271 - new_block_ids = (new_block_ids.copy(),)

275 + new_block_ids = tuple(b.copy() for b in new_block_ids)

276 load_spec = self.load_specs.pop(req_id, None)

277 request_tuple = self._unfinished_requests.get(req_id)

278 request_real = request_tuple[0] # type: ignore[index]

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

+1 -1

@@ -1497,7 +1497,7 @@ class MooncakeStoreWorker:

1497 # Single-group + PCP/DCP > 1: scale the lone group's spec.block_size to

1498 # self.block_size (= scheduler_block_size) so the coordinator's

1499 # ``block_size % hash_block_size == 0`` invariant holds.

1500 - groups = list(kv_cache_config.kv_cache_groups)

1500 + groups = list(kv_cache_config.transfer_groups)

1501 if len(groups) == 1 and groups[0].kv_cache_spec.block_size != self.block_size:

1502 g = groups[0]

1503 groups = [

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

+9 -6

@@ -90,12 +90,12 @@ class NixlBaseConnectorScheduler:

90 # Also handle unlikely SW-only model case instead of checking num_groups>1.

91 and any(

92 not isinstance(g.kv_cache_spec, FullAttentionSpec)

93 - for g in kv_cache_config.kv_cache_groups

93 + for g in kv_cache_config.transfer_groups

94 )

95 )

96 self._has_mamba = any(

97 isinstance(g.kv_cache_spec, MambaSpec)

98 - for g in kv_cache_config.kv_cache_groups

98 + for g in kv_cache_config.transfer_groups

99 )

100

101 logger.info("Initializing NIXL Scheduler %s", engine_id)

@@ -131,7 +131,7 @@ class NixlBaseConnectorScheduler:

131 (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)

132 if isinstance(g.kv_cache_spec, SlidingWindowSpec)

133 else (0, self.block_size)

134 - for g in kv_cache_config.kv_cache_groups

134 + for g in kv_cache_config.transfer_groups

135 ]

136 # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively

137 # account for boundary overlap eg window isn't fully aligned with blocks.

@@ -146,7 +146,7 @@ class NixlBaseConnectorScheduler:

146 g.kv_cache_spec.num_speculative_blocks

147 if isinstance(g.kv_cache_spec, MambaSpec)

148 else None

149 - for g in kv_cache_config.kv_cache_groups

149 + for g in kv_cache_config.transfer_groups

150 ]

151 # Only "all" mode keeps a state per block position; the other modes

152 # keep a single running state in the last non-speculative slot.

@@ -260,8 +260,11 @@ class NixlBaseConnectorScheduler:

260 for per-step partial lists (host-buffer save), where the SSM strip

261 does not apply.

262 """

263 - if len(block_ids) == 0 or not self._is_hma_required:

264 - # No blocks to clip eg Full prefix cache hit or not a hybrid model.

263 + if len(block_ids) == 0:

264 + # No blocks to clip, e.g. a full prefix cache hit.

265 + return block_ids

266 + block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids)

267 + if not self._is_hma_required:

268 return block_ids

269 # NOTE (NickLucche) This logic is currently handled at the connector level

270 # because offloading connectors might want to receive the whole sequence even

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+6 -6

@@ -300,13 +300,13 @@ class NixlBaseConnectorWorker:

300 not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager

301 and any(

302 not isinstance(g.kv_cache_spec, FullAttentionSpec)

303 - for g in kv_cache_config.kv_cache_groups

303 + for g in kv_cache_config.transfer_groups

304 )

305 )

306 self.kv_cache_config = kv_cache_config

307 self._layer_specs = {

308 layer: group.kv_cache_spec

309 - for group in kv_cache_config.kv_cache_groups

309 + for group in kv_cache_config.transfer_groups

310 for layer in group.layer_names

311 }

312

@@ -318,7 +318,7 @@ class NixlBaseConnectorWorker:

318 self._conv_decomp: MambaConvSplitInfo | None = None

319 self._has_mamba = any(

320 isinstance(g.kv_cache_spec, MambaSpec)

321 - for g in kv_cache_config.kv_cache_groups

321 + for g in kv_cache_config.transfer_groups

322 )

323 if self._has_mamba:

324 assert self._is_hma_required

@@ -554,7 +554,7 @@ class NixlBaseConnectorWorker:

554 # Unwrap UniformTypeKVCacheSpecs to get the representative spec type

555 self._group_spec_types = tuple(

556 get_representative_spec_type(g.kv_cache_spec)

557 - for g in self.kv_cache_config.kv_cache_groups

557 + for g in self.kv_cache_config.transfer_groups

558 )

559

560 # Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE

@@ -1906,7 +1906,7 @@ class NixlBaseConnectorWorker:

1906 )

1907 mamba_layers = {

1908 name

1909 - for g, group in enumerate(self.kv_cache_config.kv_cache_groups)

1909 + for g, group in enumerate(self.kv_cache_config.transfer_groups)

1910 if _is_ssm_spec(self._group_spec_types[g])

1911 for name in group.layer_names

1912 }

@@ -2361,7 +2361,7 @@ class NixlBaseConnectorWorker:

2361 return block_ids

2362 block_arange = np.arange(0, ratio).reshape(1, -1)

2363 # Mamba blocks have no logical<>physical discrepancy (block-size=1)

2364 - group_specs = self.kv_cache_config.kv_cache_groups

2364 + group_specs = self.kv_cache_config.transfer_groups

2365 physical_block_ids = []

2366 for i, group in enumerate(block_ids):

2367 spec = group_specs[i].kv_cache_spec

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

+1 -1

@@ -290,7 +290,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

290 assert (

291 len(remote_block_ids)

292 == len(local_block_ids)

293 - == len(self.kv_cache_config.kv_cache_groups)

293 + == len(self.kv_cache_config.transfer_groups)

294 )

295 local_block_ids, remote_block_ids = self._apply_prefix_caching(

296 decode_block_ids=local_block_ids,

vllm/v1/kv_cache_interface.py

+40 -1

@@ -5,10 +5,11 @@ from __future__ import annotations

6 import copy

7 from collections import Counter

8 -from collections.abc import Collection

8 +from collections.abc import Collection, Sequence

9 from dataclasses import dataclass, fields, replace

10 from enum import Enum, IntEnum

11 from fractions import Fraction

12 +from functools import cached_property

13 from math import prod

14 from typing import TYPE_CHECKING, TypeVar

@@ -1143,6 +1144,8 @@ class KVCacheGroupSpec:

1144 kv_cache_spec: KVCacheSpec

1145 # Whether this group contains EAGLE/MTP draft attention layers.

1146 is_eagle_group: bool = False

1147 + # Whether this group is part of the externally transferable KV state.

1148 + enable_kv_transfer: bool = True

1149

1150

1151 @dataclass

@@ -1168,6 +1171,42 @@ class KVCacheConfig:

1171 kv_cache_layout: str | None = None

1172 """The KV cache layout resolved by the engine core, adopted by all workers."""

1173

1174 + @cached_property

1175 + def transfer_group_ids(self) -> tuple[int, ...]:

1176 + """IDs of cache groups that participate in external KV transfer."""

1177 + return tuple(

1178 + group_id

1179 + for group_id, group in enumerate(self.kv_cache_groups)

1180 + if group.enable_kv_transfer

1181 + )

1182 +

1183 + @cached_property

1184 + def transfer_groups(self) -> tuple[KVCacheGroupSpec, ...]:

1185 + """Cache groups that participate in external KV transfer."""

1186 + return tuple(

1187 + self.kv_cache_groups[group_id] for group_id in self.transfer_group_ids

1188 + )

1189 +

1190 + @cached_property

1191 + def transfer_group_index_by_layer(self) -> dict[str, int]:

1192 + """Transfer-group tuple index for each participating layer."""

1193 + return {

1194 + layer_name: group_index

1195 + for group_index, group in enumerate(self.transfer_groups)

1196 + for layer_name in group.layer_names

1197 + }

1198 +

1199 + def select_transfer_block_ids(

1200 + self, block_ids: Sequence[list[int]]

1201 + ) -> tuple[list[int], ...]:

1202 + """Select block IDs for externally transferable cache groups."""

1203 + if len(block_ids) != len(self.kv_cache_groups):

1204 + raise ValueError(

1205 + f"Expected {len(self.kv_cache_groups)} KV cache groups, "

1206 + f"got {len(block_ids)}."

1207 + )

1208 + return tuple(block_ids[group_id] for group_id in self.transfer_group_ids)

1209 +

1210 @property

1211 def has_mamba_layers(self) -> bool:

1212 return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)