class UMBPStoreConnector(KVConnectorBase_V1):
"""vLLM connector lifecycle shared by every UMBP runtime mode."""
def __init__(
self,
vllm_config: VllmConfig,
role: KVConnectorRole,
kv_cache_config: KVCacheConfig,
) -> None:
super().__init__(vllm_config, role, kv_cache_config)
topology = RankTopology.from_vllm_config(vllm_config)
runtime_config = UMBPRuntimeConfig.from_vllm(
vllm_config
).resolve_for_rank_count(topology.rank_count)
if topology.pp_size > 1:
raise NotImplementedError(
f"{runtime_config.mode} UMBP does not support pipeline parallelism"
)
runtime = UMBPRuntimeFactory.build(runtime_config)
self._runtime = runtime
self._topology = topology
self._store_submitted = False
self._namespace: UMBPNamespace | None = None
self.connector_scheduler: UMBPStoreConnectorScheduler | None = None
self.connector_worker: UMBPStoreConnectorWorker | None = None
if role == KVConnectorRole.SCHEDULER:
self.connector_scheduler = UMBPStoreConnectorScheduler(
vllm_config, kv_cache_config, None, None, topology
)
namespace = UMBPNamespace.configured(vllm_config)
if namespace is not None:
self._bind_scheduler_namespace(namespace)
else:
self._namespace = UMBPNamespace.from_vllm_config(
vllm_config, kv_cache_config
)
layout = KVLayoutPlanner.from_kv_cache_config(kv_cache_config)
self.connector_worker = UMBPStoreConnectorWorker(
runtime.create_worker_handle(
self._namespace.value, topology, layout.describe(topology)
),
layout,
codec=self._codec(self._namespace),
)
def _codec(self, namespace: UMBPNamespace) -> BlockIdentityCodec:
topology = self._topology
return BlockIdentityCodec(
namespace=namespace,
tp_rank=topology.tp_rank,
pp_rank=topology.pp_rank,
pcp_rank=topology.pcp_rank,
dcp_rank=topology.dcp_rank,
)
def _bind_scheduler_namespace(self, namespace: UMBPNamespace) -> None:
assert self.connector_scheduler is not None
self._namespace = namespace
# Only workers see per-layer specs, so only they plan byte ranges.
layout = KVLayoutDescriptor(regions=(), topology=self._topology)
self.connector_scheduler.bind_namespace(
self._runtime.create_scheduler_handle(
namespace.value, self._topology, layout
),
self._codec(namespace),
)
def get_handshake_metadata(self) -> UMBPHandshakeMetadata | None:
if self.connector_worker is None or self._namespace is None:
return None
return UMBPHandshakeMetadata(self._namespace.value)
def set_xfer_handshake_metadata(
self, metadata: dict[int, KVConnectorHandshakeMetadata]
) -> None:
if self.connector_scheduler is None or self._namespace is not None:
return
# Reports are keyed by worker rank, but ranks may share a key (PCP
# without DCP), so require agreement rather than one per rank.
namespaces = {
report.namespace
for report in metadata.values()
if isinstance(report, UMBPHandshakeMetadata)
}
if len(namespaces) != 1:
raise ValueError(
"UMBP workers must report exactly one namespace, got "
f"{sorted(namespaces)}"
)
self._bind_scheduler_namespace(UMBPNamespace(namespaces.pop()))
def get_num_new_matched_tokens(
self, request: Request, num_computed_tokens: int
) -> tuple[int | None, bool]:
assert self.connector_scheduler is not None
return self.connector_scheduler.get_num_new_matched_tokens(
request, num_computed_tokens
)
def update_state_after_alloc(
self,
request: Request,
blocks: KVCacheBlocks,
num_external_tokens: int,
) -> None:
assert self.connector_scheduler is not None
self.connector_scheduler.update_state_after_alloc(
request, blocks, num_external_tokens
)
def bind_gpu_block_pool(self, gpu_block_pool: Any) -> None:
assert self.connector_scheduler is not None
self.connector_scheduler.bind_gpu_block_pool(gpu_block_pool)
def build_connector_meta(
self, scheduler_output: SchedulerOutput
) -> KVConnectorMetadata:
assert self.connector_scheduler is not None
return self.connector_scheduler.build_connector_meta(scheduler_output)
def request_finished(
self, request: Request, block_ids: list[int]
) -> tuple[bool, dict[str, Any] | None]:
assert self.connector_scheduler is not None
return self.connector_scheduler.request_finished(request, (block_ids,))
def update_connector_output(self, connector_output: Any) -> None:
assert self.connector_scheduler is not None
self.connector_scheduler.update_connector_output(connector_output)
def has_pending_push_work(self) -> bool:
assert self.connector_scheduler is not None
return self.connector_scheduler.has_pending_push_work()
def reset_cache(self) -> bool | None:
if self.connector_scheduler is None:
return None
return self.connector_scheduler.reset_store()
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
assert self.connector_worker is not None
self.connector_worker.register_kv_caches(kv_caches)
def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
super().bind_connector_metadata(connector_metadata)
self._store_submitted = False
def clear_connector_metadata(self) -> None:
# No-forward steps skip the model runner's wait_for_save hook.
if self.connector_worker is not None and self.has_connector_metadata():
self.wait_for_save()
super().clear_connector_metadata()
def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata) -> None:
assert self.connector_worker is not None
assert isinstance(kv_connector_metadata, UMBPConnectorMetadata)
self.connector_worker.handle_preemptions(kv_connector_metadata)
def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
del kwargs
assert self.connector_worker is not None
metadata = self._get_connector_metadata()
assert isinstance(metadata, UMBPConnectorMetadata)
self.connector_worker.start_load_kv(forward_context, metadata)
def wait_for_layer_load(self, layer_name: str) -> None:
assert self.connector_worker is not None
self.connector_worker.wait_for_layer_load(layer_name)
def save_kv_layer(
self,
layer_name: str,
kv_layer: torch.Tensor,
attn_metadata: AttentionMetadata,
**kwargs: Any,
) -> None:
# Stores are submitted once per step, from wait_for_save.
return
def wait_for_save(self) -> None:
if self._store_submitted:
return
assert self.connector_worker is not None
metadata = self._get_connector_metadata()
assert isinstance(metadata, UMBPConnectorMetadata)
self.connector_worker.enqueue_stores(metadata)
self.connector_worker.wait_for_save()
self._store_submitted = True
def get_finished(
self, finished_req_ids: set[str]
) -> tuple[set[str] | None, set[str] | None]:
assert self.connector_worker is not None
return self.connector_worker.get_finished(finished_req_ids)
def get_transfer_results(
self, finished_req_ids: set[str]
) -> KVConnectorTransferResults:
assert self.connector_worker is not None
finished_sending, finished_recving = self.connector_worker.get_finished(
finished_req_ids
)
return KVConnectorTransferResults(
finished_sending=set(finished_sending or ()),
finished_recving=set(finished_recving or ()),
)
def get_block_ids_with_load_errors(self) -> set[int]:
assert self.connector_worker is not None
return self.connector_worker.get_block_ids_with_load_errors()
def build_connector_worker_meta(self) -> UMBPConnectorWorkerMetadata:
assert self.connector_worker is not None
return self.connector_worker.build_connector_worker_meta()
def shutdown(self) -> None:
if self.connector_scheduler is not None:
self.connector_scheduler.close()
if self.connector_worker is not None:
self.connector_worker.close()