| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import TYPE_CHECKING |
| 4 | |
| 5 | from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType |
| 6 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 7 | from vllm.distributed.kv_transfer.kv_connector.v1 import ( |
| 8 | KVConnectorBase_V1, |
| 9 | KVConnectorRole, |
| 10 | ) |
| 11 | |
| 12 | if TYPE_CHECKING: |
| 13 | from vllm.config import VllmConfig |
| 14 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 15 | |
| 16 | _KV_CONNECTOR_AGENT: KVConnectorBaseType | None = None |
| 17 | |
| 18 | |
| 19 | def get_kv_transfer_group() -> KVConnectorBaseType: |
| 20 | assert _KV_CONNECTOR_AGENT is not None, ( |
| 21 | "disaggregated KV cache transfer parallel group is not initialized" |
| 22 | ) |
| 23 | return _KV_CONNECTOR_AGENT |
| 24 | |
| 25 | |
| 26 | def has_kv_transfer_group() -> bool: |
| 27 | return _KV_CONNECTOR_AGENT is not None |
| 28 | |
| 29 | |
| 30 | def is_v1_kv_transfer_group(connector: KVConnectorBaseType | None = None) -> bool: |
| 31 | """Check if the KV connector is the v1 connector. |
| 32 | If the argument is None, it will check the global KV connector |
| 33 | |
| 34 | Args: |
| 35 | connector: The KV connector to check. If None, it will check the |
| 36 | global KV connector. |
| 37 | |
| 38 | Note: |
| 39 | This function will no-longer be needed after the v1 KV connector |
| 40 | becomes the default. |
| 41 | |
| 42 | """ |
| 43 | if connector is None: |
| 44 | connector = _KV_CONNECTOR_AGENT |
| 45 | |
| 46 | if connector is None: |
| 47 | return False |
| 48 | |
| 49 | return isinstance(connector, KVConnectorBase_V1) |
| 50 | |
| 51 | |
| 52 | def _sync_engine_id_across_tp(vllm_config: "VllmConfig") -> None: |
| 53 | """Broadcast engine_id from TP rank 0 so all workers in a |
| 54 | multi-node TP group share the same value. |
| 55 | |
| 56 | When PP is enabled, also broadcast across PP ranks so all workers in |
| 57 | the same model-parallel engine share the same value. |
| 58 | """ |
| 59 | from vllm.distributed.parallel_state import ( |
| 60 | get_pp_group, |
| 61 | get_tp_group, |
| 62 | ) |
| 63 | |
| 64 | assert vllm_config.kv_transfer_config is not None |
| 65 | synced_id = get_tp_group().broadcast_object( |
| 66 | vllm_config.kv_transfer_config.engine_id, src=0 |
| 67 | ) |
| 68 | if vllm_config.parallel_config.pipeline_parallel_size > 1: |
| 69 | synced_id = get_pp_group().broadcast_object(synced_id, src=0) |
| 70 | vllm_config.kv_transfer_config.engine_id = synced_id |
| 71 | |
| 72 | |
| 73 | def ensure_kv_transfer_initialized( |
| 74 | vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig" |
| 75 | ) -> None: |
| 76 | """Initialize KV cache transfer parallel group.""" |
| 77 | global _KV_CONNECTOR_AGENT |
| 78 | |
| 79 | if _KV_CONNECTOR_AGENT is not None: |
| 80 | return |
| 81 | |
| 82 | kv_transfer_config = vllm_config.kv_transfer_config |
| 83 | if kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance: |
| 84 | # NIXL P/D requires an interleave_size equal to block_size. |
| 85 | vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config) |
| 86 | _sync_engine_id_across_tp(vllm_config) |
| 87 | |
| 88 | _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector( |
| 89 | config=vllm_config, |
| 90 | role=KVConnectorRole.WORKER, |
| 91 | kv_cache_config=kv_cache_config, |
| 92 | ) |
| 93 | |
| 94 | |
| 95 | def ensure_kv_transfer_shutdown() -> None: |
| 96 | global _KV_CONNECTOR_AGENT |
| 97 | if _KV_CONNECTOR_AGENT is not None: |
| 98 | _KV_CONNECTOR_AGENT.shutdown() |
| 99 | _KV_CONNECTOR_AGENT = None |