Skip to content

vllm.distributed.kv_transfer

Modules:

Name Description
kv_connector
kv_connector_agent

A centralized entrypoint to perform distributed KV cache transfer.
kv_lookup_buffer
kv_pipe
kv_transfer_state

KVConnectorBaseType module-attribute 

KVConnectorBaseType = Union[
    KVConnectorBase, KVConnectorBase_V1
]

__all__ module-attribute 

__all__ = [
    "get_kv_transfer_group",
    "has_kv_transfer_group",
    "is_v1_kv_transfer_group",
    "ensure_kv_transfer_initialized",
    "KVConnectorBaseType",
]

ensure_kv_transfer_initialized 

ensure_kv_transfer_initialized(
    vllm_config: VllmConfig,
) -> None

Initialize KV cache transfer parallel group.

Source code in vllm/distributed/kv_transfer/kv_transfer_state.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
    """
    Initialize KV cache transfer parallel group.
    """

    global _KV_CONNECTOR_AGENT

    if vllm_config.kv_transfer_config is None:
        return

    if (vllm_config.kv_transfer_config.is_kv_transfer_instance
            and _KV_CONNECTOR_AGENT is None):
        if envs.VLLM_USE_V1:
            _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v1(
                config=vllm_config, role=KVConnectorRole.WORKER)
        else:
            _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v0(
                rank=get_world_group().rank,
                local_rank=get_world_group().local_rank,
                config=vllm_config,
            )

get_kv_transfer_group 

get_kv_transfer_group() -> KVConnectorBaseType
Source code in vllm/distributed/kv_transfer/kv_transfer_state.py 
19
20
21
22
def get_kv_transfer_group() -> KVConnectorBaseType:
    assert _KV_CONNECTOR_AGENT is not None, (
        "disaggregated KV cache transfer parallel group is not initialized")
    return _KV_CONNECTOR_AGENT

has_kv_transfer_group 

has_kv_transfer_group() -> bool
Source code in vllm/distributed/kv_transfer/kv_transfer_state.py 
25
26
def has_kv_transfer_group() -> bool:
    return _KV_CONNECTOR_AGENT is not None

is_v1_kv_transfer_group 

is_v1_kv_transfer_group(
    connector: Optional[KVConnectorBaseType] = None,
) -> bool

Check if the KV connector is the v1 connector. If the argument is None, it will check the global KV connector

Parameters:

Name Type Description Default
connector Optional[KVConnectorBaseType]

The KV connector to check. If None, it will check the global KV connector.

 None
Note

This function will no-longer be needed after the v1 KV connector becomes the default.

Source code in vllm/distributed/kv_transfer/kv_transfer_state.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def is_v1_kv_transfer_group(
        connector: Optional[KVConnectorBaseType] = None) -> bool:
    """Check if the KV connector is the v1 connector.
    If the argument is None, it will check the global KV connector

    Args:
        connector: The KV connector to check. If None, it will check the
            global KV connector.

    Note:
        This function will no-longer be needed after the v1 KV connector
        becomes the default.
    """
    if connector is None:
        connector = _KV_CONNECTOR_AGENT

    if connector is None:
        return False

    return isinstance(connector, KVConnectorBase_V1)