Skip to content

vllm.v1.worker.gpu.kv_connector

Classes:

ActiveKVConnector

Bases: KVConnector

Source code in vllm/v1/worker/gpu/kv_connector.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class ActiveKVConnector(KVConnector):
    def __init__(
        self, vllm_config: VllmConfig, kv_caches_dict: dict[str, torch.Tensor]
    ):
        self.vllm_config = vllm_config
        self.kv_connector = get_kv_transfer_group()
        # Register kv caches with KV Connector if applicable.
        self.kv_connector.register_kv_caches(kv_caches_dict)
        self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks)

        self._pending_load_kwargs: dict[str, Any] | None = None
        self._disabled = False

    def pre_forward(
        self,
        scheduler_output: "SchedulerOutput",
        input_batch: "InputBatch | None" = None,
        attn_metadata: dict[str, Any] | None = None,
    ) -> None:
        if self._disabled:
            return

        kv_connector_metadata = scheduler_output.kv_connector_metadata
        assert kv_connector_metadata is not None
        self.kv_connector.handle_preemptions(kv_connector_metadata)
        self.kv_connector.bind_connector_metadata(kv_connector_metadata)
        self._pending_load_kwargs = self._build_load_kwargs(
            scheduler_output, input_batch, attn_metadata
        )

        if scheduler_output.has_sync_kv_loads:
            # Sync loads need to run before this step's forward.
            self._start_load_kv()
        # Otherwise start the async load after forward to keep submission
        # off the critical path.

    @staticmethod
    def _build_load_kwargs(
        scheduler_output: "SchedulerOutput",
        input_batch: "InputBatch | None",
        attn_metadata: dict[str, Any] | None,
    ) -> dict[str, Any]:
        """Keyword args for the connector's start_load_kv()."""
        if input_batch is None:
            # No forward this step.
            return {}
        load_kwargs: dict[str, Any] = dict(
            scheduler_output=scheduler_output,
            request_state_indices=input_batch.idx_mapping,
            request_ids=input_batch.req_ids,
            num_tokens=input_batch.num_tokens,
        )
        if attn_metadata is not None:
            load_kwargs["attn_metadata"] = attn_metadata
        return load_kwargs

    def _start_load_kv(self) -> None:
        load_kwargs = self._pending_load_kwargs
        assert load_kwargs is not None
        self._pending_load_kwargs = None
        # TODO: sort out KV Connectors' use of forward_context
        if is_forward_context_available():
            self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs)
        else:
            with set_forward_context(None, self.vllm_config):
                self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs)

    def finish_forward(self) -> None:
        if not self._disabled:
            self.kv_connector.finish_forward()

    def reset_capture_state(self) -> None:
        self.kv_connector.reset_capture_state()

    def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:
        if self._disabled:
            return None

        if self._pending_load_kwargs is not None:
            self._start_load_kv()

        output = KVConnectorOutput()
        self.kv_connector.wait_for_save()
        transfer_results = self.kv_connector.get_transfer_results(finished_req_ids)
        output.finished_sending = transfer_results.finished_sending or None
        output.finished_recving = transfer_results.finished_recving or None
        output.failed_recving = transfer_results.failed_recving
        output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors()
        output.kv_connector_stats = self.kv_connector.get_kv_connector_stats()
        output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events()
        output.kv_connector_worker_meta = (
            self.kv_connector.build_connector_worker_meta()
        )
        self.kv_connector.clear_connector_metadata()
        return output

    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
        if self._disabled:
            return EMPTY_MODEL_RUNNER_OUTPUT

        self.pre_forward(scheduler_output)
        self.finish_forward()
        finished_req_ids = scheduler_output.finished_req_ids
        kv_connector_output = self.post_forward(finished_req_ids)
        return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

    def set_disabled(self, disabled: bool) -> None:
        # Ensure that layer-wise connector hooks aren't called when disabled.
        kv_transfer_state._KV_CONNECTOR_AGENT = None if disabled else self.kv_connector
        self._disabled = disabled

_build_load_kwargs(scheduler_output, input_batch, attn_metadata) staticmethod

Keyword args for the connector's start_load_kv().

Source code in vllm/v1/worker/gpu/kv_connector.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
@staticmethod
def _build_load_kwargs(
    scheduler_output: "SchedulerOutput",
    input_batch: "InputBatch | None",
    attn_metadata: dict[str, Any] | None,
) -> dict[str, Any]:
    """Keyword args for the connector's start_load_kv()."""
    if input_batch is None:
        # No forward this step.
        return {}
    load_kwargs: dict[str, Any] = dict(
        scheduler_output=scheduler_output,
        request_state_indices=input_batch.idx_mapping,
        request_ids=input_batch.req_ids,
        num_tokens=input_batch.num_tokens,
    )
    if attn_metadata is not None:
        load_kwargs["attn_metadata"] = attn_metadata
    return load_kwargs

KVConnector

KVConnector interface used by GPUModelRunner.

Source code in vllm/v1/worker/gpu/kv_connector.py 
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
class KVConnector:
    """KVConnector interface used by GPUModelRunner."""

    def pre_forward(
        self,
        scheduler_output: "SchedulerOutput",
        input_batch: "InputBatch | None" = None,
        attn_metadata: dict[str, Any] | None = None,
    ) -> None:
        pass

    def finish_forward(self) -> None:
        pass

    def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:
        return None

    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
        return EMPTY_MODEL_RUNNER_OUTPUT

    def set_disabled(self, disabled: bool) -> None:
        pass

    def reset_capture_state(self) -> None:
        pass