| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import uuid |
| 5 | from dataclasses import field |
| 6 | from typing import Any, Literal, get_args |
| 7 | |
| 8 | from vllm.config.utils import config |
| 9 | from vllm.utils.hashing import safe_hash |
| 10 | |
| 11 | KVProducer = Literal["kv_producer", "kv_both"] |
| 12 | KVConsumer = Literal["kv_consumer", "kv_both"] |
| 13 | KVRole = Literal[KVProducer, KVConsumer] |
| 14 | |
| 15 | |
| 16 | def hisparse_host_pool_gib( |
| 17 | kv_transfer_config: "KVTransferConfig | None", |
| 18 | ) -> float | None: |
| 19 | if kv_transfer_config is None: |
| 20 | return None |
| 21 | if kv_transfer_config.kv_connector == "MultiConnector": |
| 22 | connectors = kv_transfer_config.kv_connector_extra_config.get("connectors", []) |
| 23 | else: |
| 24 | connectors = [ |
| 25 | { |
| 26 | "kv_connector": kv_transfer_config.kv_connector, |
| 27 | "kv_connector_extra_config": ( |
| 28 | kv_transfer_config.kv_connector_extra_config |
| 29 | ), |
| 30 | } |
| 31 | ] |
| 32 | entries = [ |
| 33 | connector.get("kv_connector_extra_config", {}) |
| 34 | for connector in connectors |
| 35 | if connector.get("kv_connector") == "HiSparseConnector" |
| 36 | ] |
| 37 | if len(entries) > 1: |
| 38 | raise ValueError("Only one HiSparseConnector may be configured") |
| 39 | if not entries: |
| 40 | return None |
| 41 | host_pool_gib = entries[0].get("host_pool_gib") |
| 42 | if host_pool_gib is None: |
| 43 | raise ValueError("HiSparseConnector requires host_pool_gib") |
| 44 | host_pool_gib = float(host_pool_gib) |
| 45 | if host_pool_gib <= 0: |
| 46 | raise ValueError("HiSparseConnector host_pool_gib must be positive") |
| 47 | return host_pool_gib |
| 48 | |
| 49 | |
| 50 | def kv_buffer_device_default_factory() -> str: |
| 51 | from vllm.platforms import current_platform |
| 52 | |
| 53 | return current_platform.device_type |
| 54 | |
| 55 | |
| 56 | @config |
| 57 | class KVTransferConfig: |
| 58 | """Configuration for distributed KV cache transfer.""" |
| 59 | |
| 60 | kv_connector: str | None = None |
| 61 | """The KV connector for vLLM to transmit KV caches between vLLM instances. |
| 62 | """ |
| 63 | |
| 64 | engine_id: str | None = None |
| 65 | """The engine id for KV transfers.""" |
| 66 | |
| 67 | kv_buffer_device: str = field(default_factory=kv_buffer_device_default_factory) |
| 68 | """The device used by kv connector to buffer the KV cache. Choices are |
| 69 | 'cuda', 'cpu' and 'xpu'.""" |
| 70 | |
| 71 | kv_buffer_size: float = 1e9 |
| 72 | """The buffer size for TorchDistributedConnector. Measured in number of |
| 73 | bytes. Recommended value: 1e9 (about 1GB).""" |
| 74 | |
| 75 | kv_role: KVRole | None = None |
| 76 | """Whether this vLLM instance produces, consumes KV cache, or both. Choices |
| 77 | are 'kv_producer', 'kv_consumer', and 'kv_both'.""" |
| 78 | |
| 79 | kv_rank: int | None = None |
| 80 | """The rank of this vLLM instance in the KV cache transfer. Typical value: |
| 81 | 0 for prefill instance, 1 for decode instance. |
| 82 | Currently only 1P1D is supported.""" |
| 83 | |
| 84 | kv_parallel_size: int = 1 |
| 85 | """The number of parallel instances for KV cache transfer.""" |
| 86 | |
| 87 | kv_ip: str = "127.0.0.1" |
| 88 | """The KV connector ip, used to build distributed connection.""" |
| 89 | |
| 90 | kv_port: int = 14579 |
| 91 | """The KV connector port, used to build distributed connection.""" |
| 92 | |
| 93 | kv_connector_extra_config: dict[str, Any] = field(default_factory=dict) |
| 94 | """any extra config that the connector may need.""" |
| 95 | |
| 96 | kv_connector_module_path: str | None = None |
| 97 | """The Python module path to dynamically load the KV connector from. |
| 98 | Only supported in V1.""" |
| 99 | |
| 100 | enable_permute_local_kv: bool = False |
| 101 | """Experiment feature flag to enable HND to NHD KV Transfer""" |
| 102 | |
| 103 | kv_load_failure_policy: Literal["recompute", "fail"] = "fail" |
| 104 | """Policy for handling KV cache load failures. |
| 105 | 'recompute': reschedule the request to recompute failed blocks |
| 106 | 'fail': immediately fail the request with an error finish reason (default)""" |
| 107 | |
| 108 | def compute_hash(self) -> str: |
| 109 | """WARNING: Whenever a new field is added to this config, |
| 110 | ensure that it is included in the factors list if |
| 111 | it affects the computation graph. |
| 112 | |
| 113 | Provide a hash that uniquely identifies all the configs |
| 114 | that affect the structure of the computation |
| 115 | graph from input ids/embeddings to the final hidden states, |
| 116 | excluding anything before input ids/embeddings and after |
| 117 | the final hidden states. |
| 118 | """ |
| 119 | # no factors to consider. |
| 120 | # this config will not affect the computation graph. |
| 121 | factors: list[Any] = [] |
| 122 | hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() |
| 123 | return hash_str |
| 124 | |
| 125 | def __post_init__(self) -> None: |
| 126 | if self.engine_id is None: |
| 127 | self.engine_id = str(uuid.uuid4()) |
| 128 | |
| 129 | if self.kv_role is not None and self.kv_role not in get_args(KVRole): |
| 130 | raise ValueError( |
| 131 | f"Unsupported kv_role: {self.kv_role}. " |
| 132 | f"Supported roles are {get_args(KVRole)}" |
| 133 | ) |
| 134 | |
| 135 | if self.kv_connector is not None and self.kv_role is None: |
| 136 | raise ValueError( |
| 137 | "Please specify kv_role when kv_connector " |
| 138 | f"is set, supported roles are {get_args(KVRole)}" |
| 139 | ) |
| 140 | |
| 141 | @property |
| 142 | def is_kv_transfer_instance(self) -> bool: |
| 143 | return self.kv_connector is not None and self.kv_role in get_args(KVRole) |
| 144 | |
| 145 | @property |
| 146 | def is_kv_producer(self) -> bool: |
| 147 | return self.kv_connector is not None and self.kv_role in get_args(KVProducer) |
| 148 | |
| 149 | @property |
| 150 | def is_kv_consumer(self) -> bool: |
| 151 | return self.kv_connector is not None and self.kv_role in get_args(KVConsumer) |
| 152 | |
| 153 | def get_from_extra_config(self, key, default) -> Any: |
| 154 | return self.kv_connector_extra_config.get(key, default) |
| 155 | |
| 156 | def has_connector(self, connector_name: str) -> bool: |
| 157 | """Whether ``connector_name`` is configured, directly or in MultiConnector.""" |
| 158 | if self.kv_connector == connector_name: |
| 159 | return True |
| 160 | return self.kv_connector == "MultiConnector" and any( |
| 161 | child.get("kv_connector") == connector_name |
| 162 | for child in self.kv_connector_extra_config.get("connectors", []) |
| 163 | ) |