Skip to content

vllm.config.kv_transfer

Classes:

KVTransferConfig

Configuration for distributed KV cache transfer.

Methods:

  • compute_hash –

    WARNING: Whenever a new field is added to this config,

  • has_connector –

    Whether connector_name is configured, directly or in MultiConnector.

  • registers_kv_cache_memory –

    Whether a connector registers the device KV cache itself with UCX or

  • uses_custom_mem_pool –

    Whether a connector allocates the KV cache from its own memory pool

Attributes:

Source code in vllm/config/kv_transfer.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
@config
class KVTransferConfig:
    """Configuration for distributed KV cache transfer."""

    kv_connector: str | None = None
    """The KV connector for vLLM to transmit KV caches between vLLM instances.
    """

    engine_id: str | None = None
    """The engine id for KV transfers."""

    kv_buffer_device: str = field(default_factory=kv_buffer_device_default_factory)
    """The device used by kv connector to buffer the KV cache. Choices are
    'cuda', 'cpu' and 'xpu'."""

    kv_role: KVRole | None = None
    """Whether this vLLM instance produces, consumes KV cache, or both. Choices
    are 'kv_producer', 'kv_consumer', and 'kv_both'. With MultiConnector, the
    roles of its connectors decide this instead."""

    kv_port: int = 14579
    """The KV connector port, used to build distributed connection."""

    kv_connector_extra_config: dict[str, Any] = field(default_factory=dict)
    """any extra config that the connector may need."""

    kv_connector_module_path: str | None = None
    """The Python module path to dynamically load the KV connector from.
    Only supported in V1."""

    enable_permute_local_kv: bool = False
    """Experiment feature flag to enable HND to NHD KV Transfer"""

    kv_load_failure_policy: Literal["recompute", "fail"] = "fail"
    """Policy for handling KV cache load failures.
    'recompute': reschedule the request to recompute failed blocks
    'fail': immediately fail the request with an error finish reason (default)"""

    def compute_hash(self) -> str:
        """WARNING: Whenever a new field is added to this config,
        ensure that it is included in the factors list if
        it affects the computation graph.

        Provide a hash that uniquely identifies all the configs
        that affect the structure of the computation
        graph from input ids/embeddings to the final hidden states,
        excluding anything before input ids/embeddings and after
        the final hidden states.
        """
        # no factors to consider.
        # this config will not affect the computation graph.
        factors: list[Any] = []
        hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
        return hash_str

    def __post_init__(self) -> None:
        if self.engine_id is None:
            self.engine_id = str(uuid.uuid4())

        if self.kv_role is not None and self.kv_role not in get_args(KVRole):
            raise ValueError(
                f"Unsupported kv_role: {self.kv_role}. "
                f"Supported roles are {get_args(KVRole)}"
            )

        if self.kv_connector is not None and self.kv_role is None:
            raise ValueError(
                "Please specify kv_role when kv_connector "
                f"is set, supported roles are {get_args(KVRole)}"
            )

    @property
    def is_kv_transfer_instance(self) -> bool:
        return self.kv_connector is not None and self.kv_role in get_args(KVRole)

    @property
    def is_kv_producer(self) -> bool:
        return any(role in get_args(KVProducer) for role in self._connector_roles())

    @property
    def is_kv_consumer(self) -> bool:
        return any(role in get_args(KVConsumer) for role in self._connector_roles())

    def _connector_roles(self) -> list[str | None]:
        """The kv_role of each connector. MultiConnector does not pass its own
        kv_role to its connectors, so only theirs count."""
        if self.kv_connector is None:
            return []
        if self.kv_connector == "MultiConnector":
            return [
                child.get("kv_role")
                for child in self.kv_connector_extra_config.get("connectors", [])
            ]
        return [self.kv_role]

    def get_from_extra_config(self, key, default) -> Any:
        return self.kv_connector_extra_config.get(key, default)

    def uses_custom_mem_pool(self) -> bool:
        """Whether a connector allocates the KV cache from its own memory pool
        (custom_mem_pool), directly or in MultiConnector."""
        configs = [self.kv_connector_extra_config]
        if self.kv_connector == "MultiConnector":
            configs += [
                child.get("kv_connector_extra_config", {})
                for child in self.kv_connector_extra_config.get("connectors", [])
            ]
        return any(config.get("custom_mem_pool") for config in configs)

    def registers_kv_cache_memory(self) -> bool:
        """Whether a connector registers the device KV cache itself with UCX or
        RDMA (NIXL, Mooncake, MoRIIO, UMBP), directly or in MultiConnector. NIXL with
        a CPU kv_buffer_device registers only its host staging buffers."""
        children = [{"kv_connector": self.kv_connector}]
        if self.kv_connector == "MultiConnector":
            children = self.kv_connector_extra_config.get("connectors", [])
        for child in children:
            name = child.get("kv_connector") or ""
            if name in _KV_CACHE_REGISTERING_CONNECTORS:
                return True
            buffer_device = child.get("kv_buffer_device", self.kv_buffer_device)
            if name.startswith("Nixl") and buffer_device != "cpu":
                return True
        return False

    def has_connector(self, connector_name: str) -> bool:
        """Whether ``connector_name`` is configured, directly or in MultiConnector."""
        if self.kv_connector == connector_name:
            return True
        return self.kv_connector == "MultiConnector" and any(
            child.get("kv_connector") == connector_name
            for child in self.kv_connector_extra_config.get("connectors", [])
        )

enable_permute_local_kv = False class-attribute instance-attribute

Experiment feature flag to enable HND to NHD KV Transfer

engine_id = None class-attribute instance-attribute

The engine id for KV transfers.

kv_buffer_device = field(default_factory=kv_buffer_device_default_factory) class-attribute instance-attribute

The device used by kv connector to buffer the KV cache. Choices are 'cuda', 'cpu' and 'xpu'.

kv_connector = None class-attribute instance-attribute

The KV connector for vLLM to transmit KV caches between vLLM instances.

kv_connector_extra_config = field(default_factory=dict) class-attribute instance-attribute

any extra config that the connector may need.

kv_connector_module_path = None class-attribute instance-attribute

The Python module path to dynamically load the KV connector from. Only supported in V1.

kv_load_failure_policy = 'fail' class-attribute instance-attribute

Policy for handling KV cache load failures. 'recompute': reschedule the request to recompute failed blocks 'fail': immediately fail the request with an error finish reason (default)

kv_port = 14579 class-attribute instance-attribute

The KV connector port, used to build distributed connection.

kv_role = None class-attribute instance-attribute

Whether this vLLM instance produces, consumes KV cache, or both. Choices are 'kv_producer', 'kv_consumer', and 'kv_both'. With MultiConnector, the roles of its connectors decide this instead.

_connector_roles()

The kv_role of each connector. MultiConnector does not pass its own kv_role to its connectors, so only theirs count.

Source code in vllm/config/kv_transfer.py 
151
152
153
154
155
156
157
158
159
160
161
def _connector_roles(self) -> list[str | None]:
    """The kv_role of each connector. MultiConnector does not pass its own
    kv_role to its connectors, so only theirs count."""
    if self.kv_connector is None:
        return []
    if self.kv_connector == "MultiConnector":
        return [
            child.get("kv_role")
            for child in self.kv_connector_extra_config.get("connectors", [])
        ]
    return [self.kv_role]

compute_hash()

WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if it affects the computation graph.

Provide a hash that uniquely identifies all the configs that affect the structure of the computation graph from input ids/embeddings to the final hidden states, excluding anything before input ids/embeddings and after the final hidden states.

Source code in vllm/config/kv_transfer.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def compute_hash(self) -> str:
    """WARNING: Whenever a new field is added to this config,
    ensure that it is included in the factors list if
    it affects the computation graph.

    Provide a hash that uniquely identifies all the configs
    that affect the structure of the computation
    graph from input ids/embeddings to the final hidden states,
    excluding anything before input ids/embeddings and after
    the final hidden states.
    """
    # no factors to consider.
    # this config will not affect the computation graph.
    factors: list[Any] = []
    hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
    return hash_str

has_connector(connector_name)

Whether connector_name is configured, directly or in MultiConnector.

Source code in vllm/config/kv_transfer.py 
193
194
195
196
197
198
199
200
def has_connector(self, connector_name: str) -> bool:
    """Whether ``connector_name`` is configured, directly or in MultiConnector."""
    if self.kv_connector == connector_name:
        return True
    return self.kv_connector == "MultiConnector" and any(
        child.get("kv_connector") == connector_name
        for child in self.kv_connector_extra_config.get("connectors", [])
    )

registers_kv_cache_memory()

Whether a connector registers the device KV cache itself with UCX or RDMA (NIXL, Mooncake, MoRIIO, UMBP), directly or in MultiConnector. NIXL with a CPU kv_buffer_device registers only its host staging buffers.

Source code in vllm/config/kv_transfer.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def registers_kv_cache_memory(self) -> bool:
    """Whether a connector registers the device KV cache itself with UCX or
    RDMA (NIXL, Mooncake, MoRIIO, UMBP), directly or in MultiConnector. NIXL with
    a CPU kv_buffer_device registers only its host staging buffers."""
    children = [{"kv_connector": self.kv_connector}]
    if self.kv_connector == "MultiConnector":
        children = self.kv_connector_extra_config.get("connectors", [])
    for child in children:
        name = child.get("kv_connector") or ""
        if name in _KV_CACHE_REGISTERING_CONNECTORS:
            return True
        buffer_device = child.get("kv_buffer_device", self.kv_buffer_device)
        if name.startswith("Nixl") and buffer_device != "cpu":
            return True
    return False

uses_custom_mem_pool()

Whether a connector allocates the KV cache from its own memory pool (custom_mem_pool), directly or in MultiConnector.

Source code in vllm/config/kv_transfer.py 
166
167
168
169
170
171
172
173
174
175
def uses_custom_mem_pool(self) -> bool:
    """Whether a connector allocates the KV cache from its own memory pool
    (custom_mem_pool), directly or in MultiConnector."""
    configs = [self.kv_connector_extra_config]
    if self.kv_connector == "MultiConnector":
        configs += [
            child.get("kv_connector_extra_config", {})
            for child in self.kv_connector_extra_config.get("connectors", [])
        ]
    return any(config.get("custom_mem_pool") for config in configs)