A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 14:00 UTC
[KVConnector]: Enable Cross-layers KV cache layout for MultiConnector (#30761)
Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com>
Kfir Toledo committed
Jan 7, 2026 at 18:59 UTC
b89443b8d9c75e6e393a320103c1b6fbd2b89c68
4 files changed
+71
-11
tests/v1/kv_connector/unit/test_multi_connector.py
+45
@@ -49,6 +49,33 @@ class MockConnector(KVConnectorBase_V1):
49
) -> KVConnectorStats | None:
50
return MockConnectorStats(data=data) if data is not None else None
51
52
+ def start_load_kv(self, forward_context, **kwargs):
53
+ pass
54
+
55
+ def wait_for_layer_load(self, layer_name):
56
+ pass
57
+
58
+ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
59
+ pass
60
+
61
+ def wait_for_save(self):
62
+ pass
63
+
64
+ def build_connector_meta(self, scheduler_output):
65
+ return None
66
+
67
+ def get_num_new_matched_tokens(self, request, num_computed_tokens):
68
+ return (0, False)
69
+
70
+ def update_state_after_alloc(self, request, blocks, num_tokens) -> None:
71
+ pass
72
+
73
+
74
+class MockCrossLayerConnector(MockConnector):
75
+ @property
76
+ def prefer_cross_layer_blocks(self) -> bool:
77
+ return True
78
+
79
80
# Register the mock connector
81
KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
@@ -601,3 +628,21 @@ class TestMultiConnectorStats:
628
# One non-empty
629
stats.data["NixlConnector"].data["transfer_duration"].append(1.0)
630
assert not stats.is_empty()
631
+
632
+
633
+class TestMultiConnectorPreferCrossLayerBlocks:
634
+ def test_all_connectors_prefer_cross_layer_blocks(self):
635
+ mc = MultiConnector.__new__(MultiConnector)
636
+ mc._connectors = [
637
+ MockCrossLayerConnector.__new__(MockCrossLayerConnector),
638
+ MockCrossLayerConnector.__new__(MockCrossLayerConnector),
639
+ ]
640
+ assert mc.prefer_cross_layer_blocks is True
641
+
642
+ def test_mixed_connectors_do_not_prefer_cross_layer_blocks(self):
643
+ mc = MultiConnector.__new__(MultiConnector)
644
+ mc._connectors = [
645
+ MockCrossLayerConnector.__new__(MockCrossLayerConnector),
646
+ MockConnector.__new__(MockConnector), # default False
647
+ ]
648
+ assert mc.prefer_cross_layer_blocks is False
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+8
-8
@@ -38,7 +38,7 @@ The class provides the following primitives:
38
import enum
39
from abc import ABC, abstractmethod
40
from collections.abc import Callable, Iterable
41
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional
41
+from typing import TYPE_CHECKING, Any, Literal, Optional
42
43
import torch
44
@@ -144,15 +144,15 @@ class KVConnectorMetadata(ABC): # noqa: B024
144
class KVConnectorBase_V1(ABC):
145
"""
146
Base class for KV connectors.
147
-
148
- Attributes:
149
- prefer_cross_layer_blocks (bool): Indicates whether this connector
150
- prefers KV blocks that hold KV data for all layers (for speeding
151
- up KV data transfers).
152
- Defaults to False.
147
"""
148
155
- prefer_cross_layer_blocks: ClassVar[bool] = False
149
+ @property
150
+ def prefer_cross_layer_blocks(self) -> bool:
151
+ """
152
+ Indicates whether this connector prefers KV blocks that hold KV data for all
153
+ layers, which can speed up KV data transfers. Defaults to False.
154
+ """
155
+ return False
156
157
def __init__(
158
self,
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+14
-1
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any
7
8
import torch
9
10
-from vllm.attention.backends.abstract import AttentionMetadata
10
+from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
11
from vllm.config import VllmConfig
12
from vllm.config.kv_transfer import KVTransferConfig
13
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
@@ -138,6 +138,12 @@ class MultiConnector(KVConnectorBase_V1):
138
# Propagated from scheduler to worker side via the connector metadata.
139
self._extra_async_saves: dict[str, int] = {}
140
141
+ @property
142
+ def prefer_cross_layer_blocks(self) -> bool:
143
+ if not self._connectors:
144
+ return False
145
+ return all(c.prefer_cross_layer_blocks for c in self._connectors)
146
+
147
@classmethod
148
def _get_connector_classes_and_configs(
149
cls, vllm_config: "VllmConfig"
@@ -164,6 +170,13 @@ class MultiConnector(KVConnectorBase_V1):
170
)
171
return ret
172
173
+ def register_cross_layers_kv_cache(
174
+ self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
175
+ ):
176
+ # Register on all connectors
177
+ for c in self._connectors:
178
+ c.register_cross_layers_kv_cache(kv_cache, attn_backend)
179
+
180
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
181
for c in self._connectors:
182
c.register_kv_caches(kv_caches)
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+4
-2
@@ -4,7 +4,7 @@ from collections import defaultdict
4
from collections.abc import Iterable
5
from dataclasses import dataclass
6
from itertools import islice
7
-from typing import Any, ClassVar
7
+from typing import Any
8
9
import torch
10
@@ -44,7 +44,9 @@ class OffloadingConnectorMetadata(KVConnectorMetadata):
44
45
46
class OffloadingConnector(KVConnectorBase_V1):
47
- prefer_cross_layer_blocks: ClassVar[bool] = True
47
+ @property
48
+ def prefer_cross_layer_blocks(self) -> bool:
49
+ return True
50
51
def __init__(
52
self,