Skip to content

vllm.v1.worker.kv_connector_model_runner_mixin

Define KV connector functionality mixin for model runners.

logger module-attribute 

logger = init_logger(__name__)

KVConnectorModelRunnerMixin

Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
 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
108
109
110
111
112
113
114
115
class KVConnectorModelRunnerMixin:

    @staticmethod
    def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
        # Update KVConnector with the KVConnector metadata forward().
        if has_kv_transfer_group():
            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)

            # Background KV cache transfers happen here.
            # These transfers are designed to be async and the requests
            # involved may be disjoint from the running requests.
            # Do this here to save a collective_rpc.
            kv_connector.start_load_kv(get_forward_context())

    @staticmethod
    def maybe_wait_for_kv_save() -> None:
        if has_kv_transfer_group():
            get_kv_transfer_group().wait_for_save()

    @staticmethod
    def get_finished_kv_transfers(
        scheduler_output: "SchedulerOutput",
    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
        if has_kv_transfer_group():
            return get_kv_transfer_group().get_finished(
                scheduler_output.finished_req_ids)
        return None, None

    @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, wait_for_save=False) as kv_connector_output:
            pass

        if (not kv_connector_output.finished_sending
                and not kv_connector_output.finished_recving):
            return EMPTY_MODEL_RUNNER_OUTPUT

        output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
        output.kv_connector_output = kv_connector_output
        return output

    @staticmethod
    def maybe_get_kv_connector_output(
        scheduler_output: "SchedulerOutput"
    ) -> AbstractContextManager[Optional[KVConnectorOutput]]:
        return KVConnectorModelRunnerMixin._get_kv_connector_output(
            scheduler_output) if has_kv_transfer_group() else nullcontext()

    # This context manager must be used within an active forward context.
    # It encapsulates the entire KV conector lifecycle within execute_model
    @staticmethod
    @contextmanager
    def _get_kv_connector_output(
        scheduler_output: "SchedulerOutput",
        wait_for_save: bool = True
    ) -> 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)

        # Background KV cache transfers happen here.
        # These transfers are designed to be async and the requests
        # involved may be disjoint from the running requests.
        # Do this here to save a collective_rpc.
        kv_connector.start_load_kv(get_forward_context())
        try:
            yield output
        finally:
            if wait_for_save:
                kv_connector.wait_for_save()

            output.finished_sending, output.finished_recving = (
                kv_connector.get_finished(scheduler_output.finished_req_ids))

            kv_connector.clear_connector_metadata()

_get_kv_connector_output staticmethod 

_get_kv_connector_output(
    scheduler_output: SchedulerOutput,
    wait_for_save: bool = True,
) -> Generator[KVConnectorOutput, None, None]
Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
 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
@staticmethod
@contextmanager
def _get_kv_connector_output(
    scheduler_output: "SchedulerOutput",
    wait_for_save: bool = True
) -> 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)

    # Background KV cache transfers happen here.
    # These transfers are designed to be async and the requests
    # involved may be disjoint from the running requests.
    # Do this here to save a collective_rpc.
    kv_connector.start_load_kv(get_forward_context())
    try:
        yield output
    finally:
        if wait_for_save:
            kv_connector.wait_for_save()

        output.finished_sending, output.finished_recving = (
            kv_connector.get_finished(scheduler_output.finished_req_ids))

        kv_connector.clear_connector_metadata()

get_finished_kv_transfers staticmethod 

get_finished_kv_transfers(
    scheduler_output: SchedulerOutput,
) -> tuple[Optional[set[str]], Optional[set[str]]]
Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
50
51
52
53
54
55
56
57
@staticmethod
def get_finished_kv_transfers(
    scheduler_output: "SchedulerOutput",
) -> tuple[Optional[set[str]], Optional[set[str]]]:
    if has_kv_transfer_group():
        return get_kv_transfer_group().get_finished(
            scheduler_output.finished_req_ids)
    return None, None

kv_connector_no_forward staticmethod 

kv_connector_no_forward(
    scheduler_output: SchedulerOutput,
    vllm_config: VllmConfig,
) -> ModelRunnerOutput
Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
@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, wait_for_save=False) as kv_connector_output:
        pass

    if (not kv_connector_output.finished_sending
            and not kv_connector_output.finished_recving):
        return EMPTY_MODEL_RUNNER_OUTPUT

    output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
    output.kv_connector_output = kv_connector_output
    return output

maybe_get_kv_connector_output staticmethod 

maybe_get_kv_connector_output(
    scheduler_output: SchedulerOutput,
) -> AbstractContextManager[Optional[KVConnectorOutput]]
Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
77
78
79
80
81
82
@staticmethod
def maybe_get_kv_connector_output(
    scheduler_output: "SchedulerOutput"
) -> AbstractContextManager[Optional[KVConnectorOutput]]:
    return KVConnectorModelRunnerMixin._get_kv_connector_output(
        scheduler_output) if has_kv_transfer_group() else nullcontext()

maybe_setup_kv_connector staticmethod 

maybe_setup_kv_connector(scheduler_output: SchedulerOutput)
Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
@staticmethod
def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
    # Update KVConnector with the KVConnector metadata forward().
    if has_kv_transfer_group():
        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)

        # Background KV cache transfers happen here.
        # These transfers are designed to be async and the requests
        # involved may be disjoint from the running requests.
        # Do this here to save a collective_rpc.
        kv_connector.start_load_kv(get_forward_context())

maybe_wait_for_kv_save staticmethod 

maybe_wait_for_kv_save() -> None
Source code in vllm/v1/worker/kv_connector_model_runner_mixin.py 
45
46
47
48
@staticmethod
def maybe_wait_for_kv_save() -> None:
    if has_kv_transfer_group():
        get_kv_transfer_group().wait_for_save()