| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ECCPUConnector — CPU encoder-cache offloading. |
| 4 | |
| 5 | A role-routed shell: one instance per process. The scheduler delegate owns the |
| 6 | mmap region and offload bookkeeping; the worker delegate owns the mmap view and |
| 7 | GPU copy plumbing. An ec_both instance reuses encoder outputs it has already |
| 8 | offloaded to CPU instead of recomputing them. |
| 9 | """ |
| 10 | |
| 11 | from typing import TYPE_CHECKING, Any |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.distributed.ec_transfer.ec_connector.base import ( |
| 16 | ECConnectorBase, |
| 17 | ECConnectorRole, |
| 18 | ECConnectorWorkerMetadata, |
| 19 | ) |
| 20 | from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( |
| 21 | ECCPUConnectorMetadata, |
| 22 | ) |
| 23 | from vllm.logger import init_logger |
| 24 | |
| 25 | if TYPE_CHECKING: |
| 26 | from vllm.config import VllmConfig |
| 27 | from vllm.v1.core.sched.output import SchedulerOutput |
| 28 | from vllm.v1.outputs import ECConnectorOutput |
| 29 | from vllm.v1.request import Request |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | |
| 34 | class ECCPUConnector(ECConnectorBase): |
| 35 | """EC connector that offloads encoder cache to a shared CPU mmap region.""" |
| 36 | |
| 37 | def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole) -> None: |
| 38 | super().__init__(vllm_config=vllm_config, role=role) |
| 39 | |
| 40 | if not vllm_config.use_v2_model_runner: |
| 41 | raise ValueError( |
| 42 | "ECCPUConnector requires the V2 model runner. " |
| 43 | "Set VLLM_USE_V2_MODEL_RUNNER=1." |
| 44 | ) |
| 45 | |
| 46 | self.connector_worker = None |
| 47 | self.connector_scheduler = None |
| 48 | |
| 49 | if role == ECConnectorRole.WORKER: |
| 50 | self.connector_worker = self._make_worker(vllm_config) |
| 51 | elif role == ECConnectorRole.SCHEDULER: |
| 52 | self.connector_scheduler = self._make_scheduler(vllm_config) |
| 53 | else: |
| 54 | raise ValueError(f"Unknown ECConnectorRole: {role}") |
| 55 | |
| 56 | # Construction seams. |
| 57 | def _make_worker(self, vllm_config: "VllmConfig"): |
| 58 | # Deferred import: the worker module touches torch/CUDA at import time |
| 59 | # via the region, so keep that cost off the scheduler path. |
| 60 | from vllm.distributed.ec_transfer.ec_connector.cpu.worker import ( |
| 61 | ECCPUWorker, |
| 62 | ) |
| 63 | |
| 64 | return ECCPUWorker(vllm_config) |
| 65 | |
| 66 | def _make_scheduler(self, vllm_config: "VllmConfig"): |
| 67 | from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ( |
| 68 | ECCPUScheduler, |
| 69 | ) |
| 70 | |
| 71 | return ECCPUScheduler(vllm_config) |
| 72 | |
| 73 | # Worker-side forwarders. |
| 74 | def start_load_caches( |
| 75 | self, encoder_cache: dict[str, torch.Tensor], **kwargs |
| 76 | ) -> None: |
| 77 | assert self.connector_worker is not None |
| 78 | metadata = self._get_connector_metadata() |
| 79 | assert isinstance(metadata, ECCPUConnectorMetadata) |
| 80 | self.connector_worker.start_load_caches( |
| 81 | encoder_cache, connector_metadata=metadata |
| 82 | ) |
| 83 | |
| 84 | def save_caches( |
| 85 | self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs |
| 86 | ) -> None: |
| 87 | assert self.connector_worker is not None |
| 88 | metadata = self._get_connector_metadata() |
| 89 | assert isinstance(metadata, ECCPUConnectorMetadata) |
| 90 | self.connector_worker.save_caches( |
| 91 | encoder_cache, mm_hash, connector_metadata=metadata |
| 92 | ) |
| 93 | |
| 94 | # Scheduler-side forwarders. |
| 95 | def has_cache_item(self, identifier: str) -> bool: |
| 96 | assert self.connector_scheduler is not None |
| 97 | return self.connector_scheduler.has_cache_item(identifier) |
| 98 | |
| 99 | def ensure_cache_available( |
| 100 | self, request: "Request", num_computed_tokens: int |
| 101 | ) -> bool: |
| 102 | assert self.connector_scheduler is not None |
| 103 | return self.connector_scheduler.ensure_cache_available( |
| 104 | request, num_computed_tokens |
| 105 | ) |
| 106 | |
| 107 | def take_unavailable_requests(self) -> set[str]: |
| 108 | assert self.connector_scheduler is not None |
| 109 | return self.connector_scheduler.take_unavailable_requests() |
| 110 | |
| 111 | def update_state_after_alloc(self, request: "Request", index: int) -> None: |
| 112 | assert self.connector_scheduler is not None |
| 113 | self.connector_scheduler.update_state_after_alloc(request, index) |
| 114 | |
| 115 | def build_connector_meta( |
| 116 | self, scheduler_output: "SchedulerOutput" |
| 117 | ) -> ECCPUConnectorMetadata: |
| 118 | assert self.connector_scheduler is not None |
| 119 | return self.connector_scheduler.build_connector_meta(scheduler_output) |
| 120 | |
| 121 | def request_finished( |
| 122 | self, request: "Request" |
| 123 | ) -> tuple[bool, "dict[str, Any] | None"]: |
| 124 | assert self.connector_scheduler is not None |
| 125 | return self.connector_scheduler.request_finished(request) |
| 126 | |
| 127 | def get_finished( |
| 128 | self, finished_req_ids: set[str] |
| 129 | ) -> tuple[set[str] | None, set[str] | None]: |
| 130 | if self.connector_worker is not None: |
| 131 | self.connector_worker.flush_saves() |
| 132 | return None, None |
| 133 | |
| 134 | def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: |
| 135 | if self.connector_worker is not None: |
| 136 | return self.connector_worker.build_connector_worker_meta() |
| 137 | return None |
| 138 | |
| 139 | def update_connector_output(self, connector_output: "ECConnectorOutput") -> None: |
| 140 | if self.connector_scheduler is not None: |
| 141 | self.connector_scheduler.update_connector_output(connector_output) |
| 142 | |
| 143 | def has_pending_push_work(self) -> bool: |
| 144 | if self.connector_scheduler is not None: |
| 145 | return self.connector_scheduler.has_pending_push_work() |
| 146 | return False |
| 147 | |
| 148 | # Shared. |
| 149 | def shutdown(self) -> None: |
| 150 | if self.connector_scheduler is not None: |
| 151 | self.connector_scheduler.shutdown() |
| 152 | if self.connector_worker is not None: |
| 153 | self.connector_worker.shutdown() |