Skip to content

vllm.v1.worker.kv_connector_model_runner_mixin

Define KV connector functionality mixin for model runners.

Classes:

KVConnectorModelRunnerMixin

Methods:

Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
 23
 24
 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
 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
class KVConnectorModelRunnerMixin:
    @staticmethod
    def kv_connector_no_forward(
        scheduler_output: "SchedulerOutput", vllm_config: VllmConfig
    ) -> ModelRunnerOutput:
        # KV send/recv even if no work to do.
        with (
            set_forward_context(None, vllm_config),
            KVConnectorModelRunnerMixin._get_kv_connector_output(
                scheduler_output
            ) as kv_connector_output,
        ):
            pass

        return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

    @staticmethod
    def maybe_get_kv_connector_output(
        scheduler_output: "SchedulerOutput",
        defer_finalize: bool = False,
    ) -> AbstractContextManager[KVConnectorOutput | None]:
        return (
            KVConnectorModelRunnerMixin._get_kv_connector_output(
                scheduler_output, defer_finalize=defer_finalize
            )
            if has_kv_transfer_group()
            else nullcontext()
        )

    @staticmethod
    def finalize_kv_connector() -> None:
        """Finalize the KV connector: wait_for_save and clear metadata.

        Call after draft model forward when defer_finalize=True was used.
        """
        if has_kv_transfer_group():
            kv_connector = get_kv_transfer_group()
            kv_connector.wait_for_save()
            kv_connector.clear_connector_metadata()

    # This context manager must be used within an active forward context.
    # It encapsulates the entire KV connector lifecycle within execute_model
    @staticmethod
    @contextmanager
    def _get_kv_connector_output(
        scheduler_output: "SchedulerOutput",
        defer_finalize: bool = False,
    ) -> Generator[KVConnectorOutput, None, None]:
        output = KVConnectorOutput()

        # Update KVConnector with the KVConnector metadata forward().
        kv_connector = get_kv_transfer_group()
        assert isinstance(kv_connector, KVConnectorBase)
        assert scheduler_output.kv_connector_metadata is not None
        kv_connector.bind_connector_metadata(scheduler_output.kv_connector_metadata)

        # Start this step's KV loads, ordered after any in-flight KV block
        # zeroing. Sync loads feed this step's forward so must precede it;
        # otherwise start (async) loads after the forward launch, keeping
        # their host-side submission cost off the critical path.
        start_after_forward = not scheduler_output.has_sync_kv_loads
        if not start_after_forward:
            kv_connector.start_load_kv(get_forward_context())
        try:
            yield output
        finally:
            if start_after_forward:
                kv_connector.start_load_kv(get_forward_context())
            if not defer_finalize:
                kv_connector.wait_for_save()

            transfer_results = kv_connector.get_transfer_results(
                scheduler_output.finished_req_ids
            )
            output.finished_sending = transfer_results.finished_sending
            output.finished_recving = transfer_results.finished_recving
            output.failed_recving = transfer_results.failed_recving
            output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors()

            output.kv_connector_stats = kv_connector.get_kv_connector_stats()
            output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events()
            output.kv_connector_worker_meta = kv_connector.build_connector_worker_meta()

            if not defer_finalize:
                kv_connector.clear_connector_metadata()

finalize_kv_connector() staticmethod

Finalize the KV connector: wait_for_save and clear metadata.

Call after draft model forward when defer_finalize=True was used.

Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
52
53
54
55
56
57
58
59
60
61
@staticmethod
def finalize_kv_connector() -> None:
    """Finalize the KV connector: wait_for_save and clear metadata.

    Call after draft model forward when defer_finalize=True was used.
    """
    if has_kv_transfer_group():
        kv_connector = get_kv_transfer_group()
        kv_connector.wait_for_save()
        kv_connector.clear_connector_metadata()