Skip to content

vllm.distributed.kv_transfer.kv_connector_agent

A centralized entrypoint to perform distributed KV cache transfer.

This implementation is a shim wrapper on two APIs exposed by kv_connector: 1. send_kv_caches_and_hidden_states 2. `recv_kv_caches_and_hidden_states

logger module-attribute 

logger = init_logger(__name__)

KVTransferAgent

A class designated for distributed KV transfer

Target use cases
  1. Disaggregated prefill
  2. Remote KV cache storage
Source code in vllm/distributed/kv_transfer/kv_connector_agent.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class KVTransferAgent:
    """
    A class designated for distributed KV transfer

    Target use cases:
        1. Disaggregated prefill
        2. Remote KV cache storage
    """

    def __init__(
        self,
        rank: int,
        local_rank: int,
        config: "VllmConfig",
    ):

        self.config = config

        if config.kv_transfer_config is None:
            raise ValueError("KVTransferConfig is not set in the VllmConfig,"
                             " cannot initialize KVConnector.")

        assert self.config.kv_transfer_config.is_kv_transfer_instance, "KV"\
            "TransferAgent should only be used when kv_connector is set."

        self.connector = KVConnectorFactory.create_connector_v0(
            rank, local_rank, config)

    def send_kv_caches_and_hidden_states(
        self,
        model_executable: torch.nn.Module,
        model_input: "ModelInputForGPUWithSamplingMetadata",
        kv_caches: list[torch.Tensor],
        hidden_or_intermediate_states: Union[torch.Tensor,
                                             IntermediateTensors],
    ) -> None:

        self.connector.send_kv_caches_and_hidden_states(
            model_executable, model_input, kv_caches,
            hidden_or_intermediate_states)

    def close(self) -> None:
        self.connector.close()

    def recv_kv_caches_and_hidden_states(
        self, model_executable: torch.nn.Module,
        model_input: "ModelInputForGPUWithSamplingMetadata",
        kv_caches: list[torch.Tensor]
    ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,
               "ModelInputForGPUWithSamplingMetadata"]:

        return self.connector.recv_kv_caches_and_hidden_states(
            model_executable, model_input, kv_caches)

config instance-attribute 

config = config

connector instance-attribute 

connector = create_connector_v0(rank, local_rank, config)

__init__ 

__init__(rank: int, local_rank: int, config: VllmConfig)
Source code in vllm/distributed/kv_transfer/kv_connector_agent.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def __init__(
    self,
    rank: int,
    local_rank: int,
    config: "VllmConfig",
):

    self.config = config

    if config.kv_transfer_config is None:
        raise ValueError("KVTransferConfig is not set in the VllmConfig,"
                         " cannot initialize KVConnector.")

    assert self.config.kv_transfer_config.is_kv_transfer_instance, "KV"\
        "TransferAgent should only be used when kv_connector is set."

    self.connector = KVConnectorFactory.create_connector_v0(
        rank, local_rank, config)

close 

close() -> None
Source code in vllm/distributed/kv_transfer/kv_connector_agent.py 
66
67
def close(self) -> None:
    self.connector.close()

recv_kv_caches_and_hidden_states 

recv_kv_caches_and_hidden_states(
    model_executable: Module,
    model_input: ModelInputForGPUWithSamplingMetadata,
    kv_caches: list[Tensor],
) -> tuple[
    Union[Tensor, IntermediateTensors],
    bool,
    ModelInputForGPUWithSamplingMetadata,
]
Source code in vllm/distributed/kv_transfer/kv_connector_agent.py 
69
70
71
72
73
74
75
76
77
def recv_kv_caches_and_hidden_states(
    self, model_executable: torch.nn.Module,
    model_input: "ModelInputForGPUWithSamplingMetadata",
    kv_caches: list[torch.Tensor]
) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,
           "ModelInputForGPUWithSamplingMetadata"]:

    return self.connector.recv_kv_caches_and_hidden_states(
        model_executable, model_input, kv_caches)

send_kv_caches_and_hidden_states 

send_kv_caches_and_hidden_states(
    model_executable: Module,
    model_input: ModelInputForGPUWithSamplingMetadata,
    kv_caches: list[Tensor],
    hidden_or_intermediate_states: Union[
        Tensor, IntermediateTensors
    ],
) -> None
Source code in vllm/distributed/kv_transfer/kv_connector_agent.py 
53
54
55
56
57
58
59
60
61
62
63
64
def send_kv_caches_and_hidden_states(
    self,
    model_executable: torch.nn.Module,
    model_input: "ModelInputForGPUWithSamplingMetadata",
    kv_caches: list[torch.Tensor],
    hidden_or_intermediate_states: Union[torch.Tensor,
                                         IntermediateTensors],
) -> None:

    self.connector.send_kv_caches_and_hidden_states(
        model_executable, model_input, kv_caches,
        hidden_or_intermediate_states)