| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import hashlib |
| 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 | |
| 10 | ECProducer = Literal["ec_producer", "ec_both"] |
| 11 | ECConsumer = Literal["ec_consumer", "ec_both"] |
| 12 | ECRole = Literal[ECProducer, ECConsumer] |
| 13 | |
| 14 | |
| 15 | @config |
| 16 | class ECTransferConfig: |
| 17 | """Configuration for distributed EC cache transfer.""" |
| 18 | |
| 19 | ec_connector: str | None = None |
| 20 | """The EC connector for vLLM to transmit EC caches between vLLM instances. |
| 21 | |
| 22 | Built-in options include ``ECExampleConnector`` (shared filesystem via |
| 23 | safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine RDMA; |
| 24 | requires ``mooncake-transfer-engine`` and matching producer/consumer |
| 25 | ``ec_connector_extra_config``; see ``mooncake_ec_connector`` module docstring). |
| 26 | Set ``cross_encoder_cache`` in Mooncake extra config to reuse shared |
| 27 | Encoder outputs from Store before encoding, retaining P2P delivery. |
| 28 | """ |
| 29 | |
| 30 | engine_id: str | None = None |
| 31 | """The engine id for EC transfers.""" |
| 32 | |
| 33 | ec_buffer_device: str | None = "cuda" |
| 34 | """The device used by ec connector to buffer the EC cache. |
| 35 | Currently only support 'cuda'.""" |
| 36 | |
| 37 | ec_buffer_size: float = 1e9 |
| 38 | """The buffer size for TorchDistributedConnector. Measured in number of |
| 39 | bytes. Recommended value: 1e9 (about 1GB).""" |
| 40 | |
| 41 | ec_role: ECRole | None = None |
| 42 | """Whether this vLLM instance produces, consumes EC cache, or both. Choices |
| 43 | are 'ec_producer', 'ec_consumer', 'ec_both'.""" |
| 44 | |
| 45 | ec_rank: int | None = None |
| 46 | """The rank of this vLLM instance in the EC cache transfer. Typical value: |
| 47 | 0 for encoder, 1 for pd instance. |
| 48 | Currently only 1P1D is supported.""" |
| 49 | |
| 50 | ec_parallel_size: int = 1 |
| 51 | """The number of parallel instances for EC cache transfer. For |
| 52 | PyNcclConnector, this should be 2.""" |
| 53 | |
| 54 | ec_ip: str = "127.0.0.1" |
| 55 | """The EC connector ip, used to build distributed connection.""" |
| 56 | |
| 57 | ec_port: int = 14579 |
| 58 | """The EC connector port, used to build distributed connection.""" |
| 59 | |
| 60 | ec_connector_extra_config: dict[str, Any] = field(default_factory=dict) |
| 61 | """any extra config that the connector may need.""" |
| 62 | |
| 63 | ec_connector_module_path: str | None = None |
| 64 | """The Python module path to dynamically load the EC connector from. |
| 65 | Only supported in V1.""" |
| 66 | |
| 67 | def compute_hash(self) -> str: |
| 68 | """WARNING: Whenever a new field is added to this config, |
| 69 | ensure that it is included in the factors list if |
| 70 | it affects the computation graph. |
| 71 | |
| 72 | Provide a hash that uniquely identifies all the configs |
| 73 | that affect the structure of the computation |
| 74 | graph from input ids/embeddings to the final hidden states, |
| 75 | excluding anything before input ids/embeddings and after |
| 76 | the final hidden states. |
| 77 | """ |
| 78 | # no factors to consider. |
| 79 | # this config will not affect the computation graph. |
| 80 | factors: list[Any] = [] |
| 81 | hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest() |
| 82 | return hash_str |
| 83 | |
| 84 | def __post_init__(self) -> None: |
| 85 | if self.engine_id is None: |
| 86 | self.engine_id = str(uuid.uuid4()) |
| 87 | |
| 88 | if self.ec_role is not None and self.ec_role not in get_args(ECRole): |
| 89 | raise ValueError( |
| 90 | f"Unsupported ec_role: {self.ec_role}. " |
| 91 | f"Supported roles are {get_args(ECRole)}" |
| 92 | ) |
| 93 | |
| 94 | if self.ec_connector is not None and self.ec_role is None: |
| 95 | raise ValueError( |
| 96 | "Please specify ec_role when ec_connector " |
| 97 | f"is set, supported roles are {get_args(ECRole)}" |
| 98 | ) |
| 99 | |
| 100 | @property |
| 101 | def is_ec_transfer_instance(self) -> bool: |
| 102 | return self.ec_connector is not None and self.ec_role in get_args(ECRole) |
| 103 | |
| 104 | @property |
| 105 | def is_ec_producer(self) -> bool: |
| 106 | return self.ec_connector is not None and self.ec_role in get_args(ECProducer) |
| 107 | |
| 108 | @property |
| 109 | def is_ec_consumer(self) -> bool: |
| 110 | return self.ec_connector is not None and self.ec_role in get_args(ECConsumer) |
| 111 | |
| 112 | @property |
| 113 | def is_encode_only(self) -> bool: |
| 114 | """Whether this instance encodes but does not run the language model. |
| 115 | |
| 116 | It allocates no KV cache either -- `GPUModelRunner.get_kv_cache_spec` |
| 117 | returns {} for it -- so it is the one role that can spend accelerator |
| 118 | time and memory on frontend work. |
| 119 | """ |
| 120 | return self.is_ec_producer and not self.is_ec_consumer |
| 121 | |
| 122 | def get_from_extra_config(self, key, default) -> Any: |
| 123 | return self.ec_connector_extra_config.get(key, default) |