From a97d19d7924a337dd0a89b71e7bcf1fdae3f9e18 Mon Sep 17 00:00:00 2001 From: CharlesWu Date: Thu, 15 Oct 2026 23:12:39 +0800 Subject: [PATCH] [UMBP][1/10] Add shared connector, runtime contracts, and request lifecycle (#61789) Signed-off-by: CharlesWu Co-authored-by: Yichao Zhu --- .../kv_connector/v1/umbp/__init__.py | 41 ++ .../kv_connector/v1/umbp/connector.py | 220 +++++++ .../kv_transfer/kv_connector/v1/umbp/data.py | 620 ++++++++++++++++++ .../kv_connector/v1/umbp/runtime/__init__.py | 13 + .../kv_connector/v1/umbp/runtime/base.py | 69 ++ .../kv_connector/v1/umbp/runtime/factory.py | 76 +++ .../kv_connector/v1/umbp/scheduler.py | 518 +++++++++++++++ .../kv_connector/v1/umbp/worker.py | 240 +++++++ 8 files changed, 1797 insertions(+) create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py new file mode 100644 index 0000000000..06eee4bfab --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .connector import UMBPStoreConnector +from .data import ( + BlockIdentityCodec, + BlockTransferPlan, + KVLayoutDescriptor, + KVLayoutPlanner, + KVRange, + KVRegion, + LoadSpec, + RankTopology, + RequestTracker, + TransferJobState, + TransferJobStatus, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, + UMBPNamespace, +) +from .runtime import UMBPRuntimeConfig, UMBPRuntimeFactory + +__all__ = [ + "UMBPStoreConnector", + "BlockIdentityCodec", + "BlockTransferPlan", + "KVLayoutPlanner", + "KVLayoutDescriptor", + "KVRange", + "KVRegion", + "LoadSpec", + "RequestTracker", + "RankTopology", + "TransferJobState", + "TransferJobStatus", + "UMBPConnectorMetadata", + "UMBPConnectorWorkerMetadata", + "UMBPNamespace", + "UMBPRuntimeConfig", + "UMBPRuntimeFactory", +] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py new file mode 100644 index 0000000000..177d9fca11 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM-facing shared UMBP connector.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, + KVConnectorMetadata, + KVConnectorRole, + KVConnectorTransferResults, +) +from vllm.forward_context import ForwardContext +from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.sched.output import SchedulerOutput + +from .data import ( + BlockIdentityCodec, + KVLayoutPlanner, + RankTopology, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, + UMBPNamespace, +) +from .runtime import UMBPRuntimeConfig, UMBPRuntimeFactory +from .scheduler import UMBPStoreConnectorScheduler +from .worker import UMBPStoreConnectorWorker + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + + +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) + namespace = UMBPNamespace.from_vllm_config(vllm_config, kv_cache_config) + codec = BlockIdentityCodec( + namespace=namespace, + tp_rank=topology.tp_rank, + pp_rank=topology.pp_rank, + pcp_rank=topology.pcp_rank, + dcp_rank=topology.dcp_rank, + ) + layout = KVLayoutPlanner.from_kv_cache_config(kv_cache_config) + layout_descriptor = layout.describe(topology) + self._runtime = runtime + self._store_submitted = False + 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, + runtime.create_scheduler_handle( + namespace.value, topology, layout_descriptor + ), + codec, + topology, + ) + else: + self.connector_worker = UMBPStoreConnectorWorker( + runtime.create_worker_handle( + namespace.value, topology, layout_descriptor + ), + layout, + codec=codec, + ) + + 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() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py new file mode 100644 index 0000000000..4fa25d4071 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py @@ -0,0 +1,620 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared data types for the UMBP KV connector. + +The types in this module deliberately do not depend on a UMBP client. They +describe vLLM's logical KV objects and the physical ranges a runtime adapter +will transfer. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Sequence +from dataclasses import dataclass, field, replace +from enum import Enum +from typing import Any, NamedTuple + +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorMetadata, + KVConnectorWorkerMetadata, +) +from vllm.v1.kv_cache_interface import KVCacheConfig + + +@dataclass(frozen=True) +class RankTopology: + """Model-parallel ranks that participate in one KV object.""" + + tp_rank: int = 0 + tp_size: int = 1 + pp_rank: int = 0 + pp_size: int = 1 + pcp_rank: int = 0 + pcp_size: int = 1 + dcp_rank: int = 0 + dcp_size: int = 1 + + def __post_init__(self) -> None: + for name in ("tp_size", "pp_size", "pcp_size", "dcp_size"): + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + for rank, size, name in ( + (self.tp_rank, self.tp_size, "tp_rank"), + (self.pp_rank, self.pp_size, "pp_rank"), + (self.pcp_rank, self.pcp_size, "pcp_rank"), + (self.dcp_rank, self.dcp_size, "dcp_rank"), + ): + if rank < 0 or rank >= size: + raise ValueError(f"{name}={rank} is outside [0, {size})") + + @classmethod + def from_vllm_config(cls, vllm_config: Any) -> RankTopology: + parallel = vllm_config.parallel_config + tp_size = getattr(parallel, "tensor_parallel_size", 1) + pp_size = getattr(parallel, "pipeline_parallel_size", 1) + pcp_size = getattr(parallel, "prefill_context_parallel_size", 1) + # An engine's workers are laid out as [PP, PCP, TP], TP innermost + # (initialize_model_parallel). + rank = getattr(parallel, "rank", 0) + tp_rank = rank % tp_size + dcp_size = getattr(parallel, "decode_context_parallel_size", 1) + return cls( + tp_rank=tp_rank, + tp_size=tp_size, + pp_rank=rank // (tp_size * pcp_size) % pp_size, + pp_size=pp_size, + pcp_rank=rank // tp_size % pcp_size, + pcp_size=pcp_size, + # DCP partitions the TP workers; it does not create a second + # Cartesian worker dimension. ParallelConfig is constructed before + # process groups, so derive this stable identity from TP rank. + dcp_rank=tp_rank % dcp_size, + dcp_size=dcp_size, + ) + + @property + def local_namespace(self) -> tuple[int, int, int, int]: + return (self.tp_rank, self.pcp_rank, self.dcp_rank, self.pp_rank) + + @property + def rank_count(self) -> int: + return len(self.all_namespaces()) + + def all_namespaces(self) -> tuple[tuple[int, int, int, int], ...]: + if self.dcp_size > 1: + if self.tp_size % self.dcp_size: + raise ValueError("dcp_size must divide tp_size") + return tuple( + (tp, pcp, tp % self.dcp_size, pp) + for pp in range(self.pp_size) + for pcp in range(self.pcp_size) + for tp in range(self.tp_size) + ) + return tuple( + (tp, pcp, dcp, pp) + for pp in range(self.pp_size) + for dcp in range(self.dcp_size) + for pcp in range(self.pcp_size) + for tp in range(self.tp_size) + ) + + +@dataclass(frozen=True) +class UMBPNamespace: + """Compatibility namespace for objects stored outside the engine.""" + + value: str + + @classmethod + def from_vllm_config( + cls, + vllm_config: Any, + kv_cache_config: KVCacheConfig, + ) -> UMBPNamespace: + extra = vllm_config.kv_transfer_config.kv_connector_extra_config + configured = extra.get("key_namespace", "auto") + if configured != "auto": + if not isinstance(configured, str) or not configured: + raise ValueError("key_namespace must be a non-empty string or 'auto'") + return cls(configured) + + model_config = vllm_config.model_config + parallel_config = vllm_config.parallel_config + # Draft-model KV (e.g. DSpark/EAGLE layers) is stored alongside the + # target's, so a different drafter must not reuse it. + speculative_config = getattr(vllm_config, "speculative_config", None) + draft_model_config = getattr(speculative_config, "draft_model_config", None) + speculative = ( + None + if speculative_config is None + else { + "method": getattr(speculative_config, "method", None), + "model": getattr(draft_model_config, "model", None), + "revision": getattr(draft_model_config, "revision", None), + } + ) + fields = { + "model": getattr(model_config, "model", ""), + "revision": getattr(model_config, "revision", None), + "block_size": vllm_config.cache_config.block_size, + "cache_layout": getattr(kv_cache_config, "kv_cache_layout", None), + "groups": [ + { + "layers": group.layer_names, + "block_size": group.kv_cache_spec.block_size, + "kind": type(group.kv_cache_spec).__name__, + "dtype": str(getattr(group.kv_cache_spec, "dtype", None)), + } + for group in kv_cache_config.kv_cache_groups + ], + "tp": parallel_config.tensor_parallel_size, + "pp": parallel_config.pipeline_parallel_size, + "pcp": getattr(parallel_config, "prefill_context_parallel_size", 1), + "dcp": getattr(parallel_config, "decode_context_parallel_size", 1), + "layout_version": 1, + } + if speculative is not None: + # Added only when set, so namespaces without a drafter are unchanged. + fields["speculative"] = speculative + payload = json.dumps(fields, sort_keys=True, separators=(",", ":")).encode() + return cls(hashlib.sha256(payload).hexdigest()[:32]) + + +@dataclass(frozen=True) +class BlockIdentityCodec: + """Build stable rank-local keys from vLLM block hashes.""" + + namespace: UMBPNamespace + tp_rank: int = 0 + pp_rank: int = 0 + pcp_rank: int = 0 + dcp_rank: int = 0 + + def key(self, block_hash: bytes, group_id: int = 0) -> str: + return self.key_for_namespace( + block_hash, + group_id, + (self.tp_rank, self.pcp_rank, self.dcp_rank, self.pp_rank), + ) + + def key_for_namespace( + self, + block_hash: bytes, + group_id: int, + namespace: tuple[int, int, int, int], + ) -> str: + if not isinstance(block_hash, bytes) or not block_hash: + raise ValueError("block_hash must be non-empty bytes") + if group_id < 0: + raise ValueError("group_id must be non-negative") + tp_rank, pcp_rank, dcp_rank, pp_rank = namespace + return ( + f"umbp:vllm:v1:{self.namespace.value}:" + f"tp{tp_rank}:pcp{pcp_rank}:dcp{dcp_rank}:" + f"pp{pp_rank}:g{group_id}:" + f"{block_hash.hex()}" + ) + + def keys_for_topology( + self, + block_hash: bytes, + topology: RankTopology, + group_ids: Sequence[int], + ) -> tuple[str, ...]: + return tuple( + self.key_for_namespace(block_hash, group_id, namespace) + for namespace in topology.all_namespaces() + for group_id in group_ids + ) + + +@dataclass(frozen=True) +class KVRegion: + """One deterministic region in a rank-local KV object.""" + + layer_name: str + group_id: int + block_stride: int + block_bytes: int + object_offset: int + block_size: int = 1 + + def __post_init__(self) -> None: + if self.block_stride <= 0 or self.block_bytes <= 0: + raise ValueError("KV region sizes must be positive") + if self.block_size <= 0: + raise ValueError("KV region block_size must be positive") + if self.block_bytes > self.block_stride: + raise ValueError( + f"KV region bytes ({self.block_bytes}) exceed its stride " + f"({self.block_stride}) for {self.layer_name!r}" + ) + + +@dataclass(frozen=True) +class KVLayoutDescriptor: + """Backend-neutral description of one rank-local KV object layout.""" + + regions: tuple[KVRegion, ...] + topology: RankTopology + layout_format: str = "rank_local" + + def __post_init__(self) -> None: + if not self.layout_format: + raise ValueError("layout_format must not be empty") + names = [region.layer_name for region in self.regions] + if len(names) != len(set(names)): + raise ValueError("KV layout contains duplicate layer regions") + + @property + def object_size(self) -> int: + return sum(region.block_bytes for region in self.regions) + + +class KVLayoutPlanner: + """Create a deterministic all-layer object layout from KVCacheConfig.""" + + def __init__(self, regions: Sequence[KVRegion]) -> None: + self.regions = tuple(regions) + self._regions_by_group = { + group_id: tuple(r for r in self.regions if r.group_id == group_id) + for group_id in {r.group_id for r in self.regions} + } + self._base_addresses: dict[str, int] = {} + self._physical_strides: dict[str, int] = {} + + @classmethod + def from_kv_cache_config(cls, config: KVCacheConfig) -> KVLayoutPlanner: + group_by_layer = { + layer: group_id + for group_id, group in enumerate(config.kv_cache_groups) + for layer in group.layer_names + if group.enable_kv_transfer + } + regions: list[KVRegion] = [] + offset = 0 + for tensor in config.kv_cache_tensors: + for layer_name in tensor.layers: + group_id = group_by_layer.get(layer_name) + if group_id is None: + continue + spec = config.kv_cache_groups[group_id].kv_cache_spec + block_bytes = spec.page_size_bytes + regions.append( + KVRegion( + layer_name=layer_name, + group_id=group_id, + block_stride=tensor.block_stride, + block_bytes=block_bytes, + object_offset=offset, + block_size=spec.block_size, + ) + ) + offset += block_bytes + return cls(regions) + + @property + def object_size(self) -> int: + return sum(region.block_bytes for region in self.regions) + + def describe( + self, + topology: RankTopology, + layout_format: str = "rank_local", + ) -> KVLayoutDescriptor: + return KVLayoutDescriptor( + regions=self.regions, + topology=topology, + layout_format=layout_format, + ) + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + """Capture stable base addresses for the planned KV regions.""" + required = {region.layer_name for region in self.regions} + missing = required - kv_caches.keys() + if missing: + raise ValueError(f"missing KV cache layers: {sorted(missing)}") + addresses: dict[str, int] = {} + for layer_name in required: + cache = kv_caches[layer_name] + if cache.ndim == 0: + raise ValueError(f"KV cache {layer_name!r} must be a tensor") + block_stride = cache.stride(0) * cache.element_size() + addresses[layer_name] = cache.data_ptr() + for region in self.regions: + if region.layer_name != layer_name: + continue + if ( + block_stride > region.block_stride + or region.block_stride % block_stride != 0 + or ( + region.block_bytes > block_stride + and region.block_bytes % block_stride != 0 + ) + ): + raise ValueError( + f"KV block stride mismatch for {layer_name!r}: " + f"config={region.block_stride}, tensor={block_stride}" + ) + self._physical_strides[layer_name] = block_stride + self._base_addresses = addresses + + def plan_registered_block( + self, + key: str, + block_id: int, + *, + group_id: int | None = None, + request_id: str | None = None, + ) -> BlockTransferPlan: + if not self._base_addresses: + raise RuntimeError("KV caches must be registered before planning") + return self.plan_for_block( + key, + block_id, + self._base_addresses, + group_id=group_id, + request_id=request_id, + ) + + def plan_for_block( + self, + key: str, + block_id: int, + base_addresses: dict[str, int], + *, + group_id: int | None = None, + request_id: str | None = None, + ) -> BlockTransferPlan: + """Lower one GPU block into a deterministic scatter/gather plan.""" + return BlockTransferPlan( + key=key, + block_id=block_id, + ranges=self._ranges_for_block(block_id, base_addresses, group_id), + request_id=request_id, + ) + + def materialize(self, plan: BlockTransferPlan) -> BlockTransferPlan: + if not self._base_addresses: + raise RuntimeError("KV caches must be registered before planning") + return replace( + plan, + ranges=self._ranges_for_block( + plan.block_id, self._base_addresses, plan.group_id + ), + ) + + def _ranges_for_block( + self, + block_id: int, + base_addresses: dict[str, int], + group_id: int | None, + ) -> tuple[KVRange, ...]: + if block_id < 0: + raise ValueError("block_id must be non-negative") + regions = ( + self.regions + if group_id is None + else self._regions_by_group.get(group_id, ()) + ) + if group_id is not None and not regions: + raise ValueError(f"KV layout does not contain cache group {group_id}") + ranges: list[KVRange] = [] + object_offset = 0 + for region in regions: + try: + base_address = base_addresses[region.layer_name] + except KeyError as exc: + raise ValueError( + f"missing base address for layer {region.layer_name!r}" + ) from exc + physical_stride = self._physical_strides.get( + region.layer_name, region.block_stride + ) + if region.block_bytes <= physical_stride: + ranges.append( + KVRange( + layer_name=region.layer_name, + group_id=region.group_id, + block_id=block_id, + base_address=base_address + block_id * physical_stride, + stride=physical_stride, + length=region.block_bytes, + object_offset=( + object_offset + if group_id is not None + else region.object_offset + ), + ) + ) + object_offset += region.block_bytes + continue + physical_parts = max(1, region.block_bytes // physical_stride) + for part in range(physical_parts): + part_start = part * physical_stride + part_end = min(part_start + physical_stride, region.block_bytes) + physical_block_id = block_id * physical_parts + part + ranges.append( + KVRange( + layer_name=region.layer_name, + group_id=region.group_id, + block_id=block_id, + base_address=base_address + physical_block_id * physical_stride, + stride=physical_stride, + length=part_end - part_start, + object_offset=( + object_offset + part_start + if group_id is not None + else region.object_offset + part * physical_stride + ), + ) + ) + object_offset += region.block_bytes + return tuple(ranges) + + +class KVRange(NamedTuple): + """A source/destination range in a runtime transfer.""" + + layer_name: str + group_id: int + block_id: int + base_address: int + stride: int + length: int + object_offset: int + + +@dataclass(frozen=True) +class BlockTransferPlan: + """One complete logical block transfer.""" + + key: str + block_id: int + ranges: tuple[KVRange, ...] = () + request_id: str | None = None + group_id: int | None = None + + +@dataclass +class LoadSpec: + """Scheduler decision for one externally loaded prefix.""" + + local_tokens: int + external_tokens: int + + @property + def num_tokens_to_load(self) -> int: + return max(self.external_tokens - self.local_tokens, 0) + + +@dataclass +class RequestTracker: + """Core request state shared by every runtime mode.""" + + saved_tokens: int = 0 + block_ids: tuple[list[int], ...] = () + load_spec: LoadSpec | None = None + + def reset(self) -> None: + self.saved_tokens = 0 + self.block_ids = () + self.load_spec = None + + def update_blocks(self, block_ids: tuple[list[int], ...]) -> None: + if not self.block_ids: + self.block_ids = tuple(list(group) for group in block_ids) + return + if len(self.block_ids) != len(block_ids): + raise ValueError("KV cache group count changed for a request") + for current, new in zip(self.block_ids, block_ids, strict=True): + if new[: len(current)] == current: + current.extend(new[len(current) :]) + elif not new or current[-len(new) :] == new: + continue + else: + current.extend(new) + + def replace_blocks(self, block_ids: tuple[list[int], ...]) -> None: + """Replace stale pre-preemption tables with resumed full tables.""" + self.block_ids = tuple(list(group) for group in block_ids) + self.saved_tokens = 0 + + def mark_saved(self, token_count: int, block_size: int) -> int: + """Advance the save watermark only to a complete block boundary.""" + complete_tokens = token_count // block_size * block_size + if complete_tokens < self.saved_tokens: + return self.saved_tokens + self.saved_tokens = complete_tokens + return complete_tokens + + +class TransferJobStatus(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + + +@dataclass +class TransferJobState: + """Per-key completion state; one failed key does not fail its siblings.""" + + plans: tuple[BlockTransferPlan, ...] + status: TransferJobStatus = TransferJobStatus.PENDING + completed_keys: set[str] = field(default_factory=set) + failed_keys: set[str] = field(default_factory=set) + error: str | None = None + + @property + def keys(self) -> Sequence[str]: + return tuple(plan.key for plan in self.plans) + + def start(self) -> None: + if self.status != TransferJobStatus.PENDING: + raise RuntimeError(f"cannot start job in state {self.status}") + self.status = TransferJobStatus.RUNNING + + def complete(self, keys: Sequence[str] = ()) -> None: + self.completed_keys.update(keys or self.keys) + self._finish_if_done() + + def fail(self, keys: Sequence[str], error: str | None = None) -> None: + self.failed_keys.update(keys) + self.error = error + self._finish_if_done() + + @property + def failed_block_ids(self) -> set[int]: + return {plan.block_id for plan in self.plans if plan.key in self.failed_keys} + + def _finish_if_done(self) -> None: + keys = set(self.keys) + if self.completed_keys | self.failed_keys >= keys: + self.status = ( + TransferJobStatus.FAILED + if self.failed_keys + else TransferJobStatus.COMPLETED + ) + + +@dataclass +class UMBPConnectorMetadata(KVConnectorMetadata): + """Scheduler-to-worker metadata for one engine step.""" + + async_load: bool = False + store_plans: list[BlockTransferPlan] = field(default_factory=list) + load_requests: dict[str, list[BlockTransferPlan]] = field(default_factory=dict) + store_requests: dict[str, list[BlockTransferPlan]] = field(default_factory=dict) + preempted_request_ids: set[str] = field(default_factory=set) + store_event: int = -1 + + +@dataclass +class StoreEventResult: + """Number of workers that finished one store event.""" + + completed_workers: int = 0 + + def merge(self, other: StoreEventResult) -> None: + self.completed_workers += other.completed_workers + + +@dataclass +class UMBPConnectorWorkerMetadata(KVConnectorWorkerMetadata): + """Worker completion metadata aggregated across ranks.""" + + store_events: dict[int, StoreEventResult] = field(default_factory=dict) + + def aggregate( + self, other: KVConnectorWorkerMetadata + ) -> UMBPConnectorWorkerMetadata: + if not isinstance(other, UMBPConnectorWorkerMetadata): + raise TypeError("cannot aggregate incompatible UMBP worker metadata") + for event, result in other.store_events.items(): + self.store_events.setdefault(event, StoreEventResult()).merge(result) + return self diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py new file mode 100644 index 0000000000..7b2ce33ebf --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .base import IUMBPRuntime, UMBPSchedulerHandle, UMBPWorkerHandle +from .factory import UMBPRuntimeConfig, UMBPRuntimeFactory + +__all__ = [ + "IUMBPRuntime", + "UMBPSchedulerHandle", + "UMBPWorkerHandle", + "UMBPRuntimeConfig", + "UMBPRuntimeFactory", +] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py new file mode 100644 index 0000000000..22277492d3 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Backend-neutral runtime contract for UMBP.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +import torch + +from ..data import ( + BlockTransferPlan, + KVLayoutDescriptor, + RankTopology, + TransferJobState, +) + + +class UMBPSchedulerHandle(Protocol): + def lookup(self, keys: Sequence[str]) -> Sequence[bool]: + """Return authoritative existence for each key.""" + + def clear(self) -> bool: + """Clear all objects visible to this scheduler handle.""" + + def close(self) -> None: + """Release scheduler-side resources.""" + + +class UMBPWorkerHandle(Protocol): + def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None: + """Register each unique KV allocation with the runtime.""" + + def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: + """Start loading the requested plans.""" + + def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: + """Start storing the requested plans.""" + + def wait(self, job: TransferJobState) -> TransferJobState: + """Block until the job is final and its buffers are safe to reuse.""" + + def poll(self, job: TransferJobState) -> TransferJobState | None: + """Return a finished job without blocking, or None if still pending.""" + + def publish(self, job: TransferJobState) -> None: + """Accept a completed store; a runtime may defer visibility until then.""" + + def close(self) -> None: + """Release worker-side resources.""" + + +class IUMBPRuntime(Protocol): + def create_scheduler_handle( + self, + namespace: str, + topology: RankTopology, + layout: KVLayoutDescriptor, + ) -> UMBPSchedulerHandle: + """Create a scheduler-side handle.""" + + def create_worker_handle( + self, + namespace: str, + topology: RankTopology, + layout: KVLayoutDescriptor, + ) -> UMBPWorkerHandle: + """Create a worker-side handle.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py new file mode 100644 index 0000000000..b62c2587fc --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Runtime configuration and adapter factory for UMBP.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +from vllm.config import VllmConfig + +from .base import IUMBPRuntime + + +@dataclass(frozen=True) +class UMBPRuntimeConfig: + """Shared connector options passed to the selected runtime adapter.""" + + mode: str + options: dict[str, Any] + rank_count: int = 1 + dp_index: int = 0 + + @classmethod + def from_vllm(cls, vllm_config: VllmConfig) -> UMBPRuntimeConfig: + transfer_config = vllm_config.kv_transfer_config + if transfer_config is None: + raise ValueError("UMBP requires kv_transfer_config") + options = dict(transfer_config.kv_connector_extra_config) + mode = options.get("mode", "embedded") + if mode not in {"embedded"}: + raise ValueError(f"unknown UMBP mode: {mode!r}") + namespace = options.get("key_namespace", "auto") + if namespace != "auto" and (not isinstance(namespace, str) or not namespace): + raise ValueError("key_namespace must be a non-empty string or 'auto'") + return cls( + mode=mode, + options=options, + dp_index=int( + getattr(vllm_config.parallel_config, "data_parallel_index", 0) or 0 + ), + ) + + def resolve_for_rank_count(self, rank_count: int) -> UMBPRuntimeConfig: + """Pass topology to the adapter without interpreting its capacity options.""" + if rank_count <= 0: + raise ValueError("rank_count must be positive") + return replace(self, rank_count=rank_count) + + +RuntimeBuilder = Callable[[UMBPRuntimeConfig], IUMBPRuntime] + + +class UMBPRuntimeFactory: + """Registry used by UMBP adapters without importing MORI in vLLM.""" + + _builders: dict[str, RuntimeBuilder] = {} + + @classmethod + def register(cls, mode: str, builder: RuntimeBuilder) -> None: + if mode not in {"embedded"}: + raise ValueError(f"unknown UMBP mode: {mode!r}") + if mode in cls._builders: + raise ValueError(f"UMBP runtime mode {mode!r} is already registered") + cls._builders[mode] = builder + + @classmethod + def build(cls, config: UMBPRuntimeConfig) -> IUMBPRuntime: + try: + builder = cls._builders[config.mode] + except KeyError as exc: + raise RuntimeError( + f"no UMBP runtime adapter is registered for mode {config.mode!r}" + ) from exc + return builder(config) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py new file mode 100644 index 0000000000..8ca17a99c1 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared scheduler-side UMBP connector behavior.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Any + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +from vllm.logger import init_logger +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.kv_cache_utils import ( + resolve_dcp_kv_block_size, + resolve_kv_cache_block_sizes, +) +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, +) +from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.request import Request + +from .data import ( + BlockIdentityCodec, + BlockTransferPlan, + LoadSpec, + RankTopology, + RequestTracker, + StoreEventResult, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, +) +from .runtime import UMBPSchedulerHandle + +logger = init_logger(__name__) + + +@lru_cache(maxsize=8192) +def _rank_keys( + codec: BlockIdentityCodec, + topology: RankTopology, + block_hash: bytes, + group_id: int, +) -> tuple[str, ...]: + """Reuse immutable object identities, not runtime residency.""" + return codec.keys_for_topology(block_hash, topology, (group_id,)) + + +@dataclass +class _StoreEvent: + plans: tuple[BlockTransferPlan, ...] + result: StoreEventResult = field(default_factory=StoreEventResult) + + +class UMBPStoreConnectorScheduler: + """Own vLLM prefix matching while the runtime owns lookup transport.""" + + def __init__( + self, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + runtime: UMBPSchedulerHandle, + codec: BlockIdentityCodec, + topology: RankTopology | None = None, + ) -> None: + self.block_size, resolved_hash_block_size = resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) + self.kv_cache_config = kv_cache_config + dcp_size = vllm_config.parallel_config.decode_context_parallel_size + self.group_block_sizes = { + group_id: resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_size) + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + } + self.runtime = runtime + self.codec = codec + self.topology = topology or RankTopology() + transfer_config = vllm_config.kv_transfer_config + assert transfer_config is not None + extra = transfer_config.kv_connector_extra_config + self.load_async = bool(extra.get("load_async", True)) + self.enable_lookup = bool(extra.get("enable_lookup", True)) + self.save_decode_cache = bool(extra.get("save_decode_cache", False)) + # A kv_consumer only restores from the pool, as with other store + # connectors; it never offloads its own KV. + self.store_enabled = getattr(transfer_config, "kv_role", None) != "kv_consumer" + # Request.block_hashes are computed at this size; keys index into them. + self.hash_block_size = resolved_hash_block_size + speculative_config = getattr(vllm_config, "speculative_config", None) + self.use_eagle = bool( + speculative_config is not None and speculative_config.use_eagle_block_drop() + ) + self._pending_loads: dict[str, list[BlockTransferPlan]] = {} + self._load_specs: dict[str, LoadSpec] = {} + self._lookup_units = { + group_id: self.group_block_sizes[group_id] + for group_id in kv_cache_config.prefix_cacheable_group_ids + } + self._requests: dict[str, Request] = {} + self._request_trackers: dict[str, RequestTracker] = {} + self._gpu_block_pool: BlockPool | None = None + self._store_events: dict[int, _StoreEvent] = {} + self._num_workers = getattr(vllm_config.parallel_config, "world_size", 1) + self._store_event_counter = 0 + + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: + self._gpu_block_pool = gpu_block_pool + + def get_num_new_matched_tokens( + self, request: Request, num_computed_tokens: int + ) -> tuple[int | None, bool]: + if not self.enable_lookup: + return 0, False + hashes = request.block_hashes + if not hashes or request.num_tokens < self.block_size: + return 0, False + if num_computed_tokens % self.block_size != 0: + return 0, False + group_ids = self.kv_cache_config.prefix_cacheable_group_ids + # vLLM must execute at least the final prompt token. Returning a hit + # for the entire prompt leaves the scheduler with no new token to run. + max_external_token = request.num_tokens - 1 + if max_external_token <= num_computed_tokens: + self._load_specs.pop(request.request_id, None) + return 0, False + logical_objects_by_group, keys = self._lookup_objects( + request, num_computed_tokens, hashes + ) + try: + hits = list(self.runtime.lookup(keys)) + except Exception as exc: + logger.debug("UMBP lookup failed request=%s: %s", request.request_id, exc) + return 0, False + if len(hits) != len(keys): + logger.debug( + "UMBP lookup returned an invalid result length request=%s", + request.request_id, + ) + return 0, False + per_rank = self.topology.rank_count + offset = 0 + logical_hits_by_group: dict[int, list[bool]] = {} + for group_id in group_ids: + group_size = len(logical_objects_by_group[group_id]) * per_rank + group_hits = hits[offset : offset + group_size] + offset += group_size + logical_hits_by_group[group_id] = ( + group_hits + if per_rank == 1 + else [ + all(group_hits[index : index + per_rank]) + for index in range(0, len(group_hits), per_rank) + ] + ) + group_id = group_ids[0] + matched_blocks = 0 + for hit in logical_hits_by_group[group_id]: + if not hit: + break + matched_blocks += 1 + if self.use_eagle: + # As core's coordinator: the drafter recomputes the last block. + matched_blocks = max(0, matched_blocks - 1) + need_to_load = matched_blocks * self._lookup_units[group_id] + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "UMBP lookup result request=%s keys=%s group_hits=%s load_tokens=%d", + request.request_id, + keys, + logical_hits_by_group, + need_to_load, + ) + matched_tokens = num_computed_tokens + need_to_load + if need_to_load <= 0: + return 0, False + self._load_specs[request.request_id] = LoadSpec( + local_tokens=num_computed_tokens, + external_tokens=matched_tokens, + ) + return need_to_load, self.load_async + + def _lookup_objects( + self, request: Request, local_tokens: int, hashes: Sequence[bytes] + ) -> tuple[dict[int, list[tuple[bytes, tuple[str, ...]]]], list[str]]: + objects: dict[int, list[tuple[bytes, tuple[str, ...]]]] = {} + keys: list[str] = [] + for group_id, unit in self._lookup_units.items(): + group_objects = [] + for token_end in range(local_tokens + unit, request.num_tokens, unit): + block_hash = self._object_hash_at_token_end(hashes, token_end) + rank_keys = _rank_keys(self.codec, self.topology, block_hash, group_id) + group_objects.append((block_hash, rank_keys)) + keys.extend(rank_keys) + objects[group_id] = group_objects + return objects, keys + + def update_state_after_alloc( + self, + request: Request, + blocks: KVCacheBlocks, + num_external_tokens: int, + ) -> None: + self._requests[request.request_id] = request + block_groups = blocks.get_block_ids( + group_ids=self.kv_cache_config.prefix_cacheable_group_ids + ) + if not block_groups: + return + tracker = self._tracker_for_request(request.request_id) + # Core passes the request's whole table on every admission. A request + # whose load failed is recomputed into newly allocated blocks. + tracker.replace_blocks(tuple(list(group) for group in block_groups)) + if num_external_tokens <= 0: + self._pending_loads.pop(request.request_id, None) + self._load_specs.pop(request.request_id, None) + return + spec = self._load_specs.get(request.request_id) + if spec is not None: + tracker.load_spec = spec + # A request admitted without its load would wait for it forever. + assert num_external_tokens == spec.num_tokens_to_load, ( + f"UMBP matched {spec.num_tokens_to_load} external tokens for " + f"request {request.request_id} but {num_external_tokens} were " + "allocated" + ) + plans = self._load_plans_for_external_tokens( + request, + tracker, + block_groups, + num_external_tokens, + ) + self._pending_loads[request.request_id] = plans + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + meta = UMBPConnectorMetadata(async_load=self.load_async) + for request_id in scheduler_output.finished_req_ids: + self._pending_loads.pop(request_id, None) + self._load_specs.pop(request_id, None) + self._request_trackers.pop(request_id, None) + self._requests.pop(request_id, None) + for request in scheduler_output.scheduled_new_reqs: + load_plans = self._pending_loads.pop(request.req_id, []) + self._load_specs.pop(request.req_id, None) + if load_plans: + meta.load_requests[request.req_id] = load_plans + else: + tracker = self._tracker_for_request(request.req_id) + total_tokens = ( + request.num_computed_tokens + + scheduler_output.num_scheduled_tokens[request.req_id] + ) + tracked_request = self._requests.get(request.req_id) + selected_block_ids = tuple( + request.block_ids[group_id] + for group_id in (self.kv_cache_config.prefix_cacheable_group_ids) + ) + store_plans = ( + [] + if not self.store_enabled or tracked_request is None + else self._store_plans( + tracked_request, + tracker, + total_tokens, + block_ids_override=selected_block_ids, + ) + ) + logger.debug( + "UMBP store planning request=%s tokens=%d plans=%d keys=%s", + request.req_id, + total_tokens, + len(store_plans), + tuple(plan.key for plan in store_plans), + ) + meta.store_plans.extend(store_plans) + if store_plans: + meta.store_requests[request.req_id] = store_plans + + block_state = getattr(scheduler_output, "kv_connector_block_state", None) + cached_reqs = scheduler_output.scheduled_cached_reqs + resumed_ids: set[str] = set(getattr(cached_reqs, "resumed_req_ids", ())) + for request_id in cached_reqs.req_ids: + load_plans = self._pending_loads.pop(request_id, []) + self._load_specs.pop(request_id, None) + if load_plans: + meta.load_requests[request_id] = load_plans + elif self.store_enabled: + cached_request = self._requests.get(request_id) + cached_tracker = self._request_trackers.get(request_id) + if cached_request is not None and cached_tracker is not None: + request_index = cached_reqs.req_ids.index(request_id) + num_computed_tokens = cached_reqs.num_computed_tokens[request_index] + # num_tokens counts sampled tokens, so a decode step is still + # short of it; only the prompt bounds the prefill. + is_prefill = num_computed_tokens < cached_request.num_prompt_tokens + if not is_prefill and not self.save_decode_cache: + continue + new_block_ids = cached_reqs.new_block_ids[request_index] + if new_block_ids: + selected = tuple( + new_block_ids[group_id] + for group_id in ( + self.kv_cache_config.prefix_cacheable_group_ids + ) + ) + if request_id in resumed_ids: + cached_tracker.replace_blocks(selected) + else: + cached_tracker.update_blocks(selected) + # The appended table keeps blocks core has since freed, such + # as sliding-window blocks out of the window, which a store + # would read; core's current table does not. + current = ( + block_state.get_block_ids(request_id) + if block_state is not None + else None + ) + if current is not None: + cached_tracker.block_ids = tuple( + list(current[group_id]) + for group_id in ( + self.kv_cache_config.prefix_cacheable_group_ids + ) + ) + total_tokens = ( + num_computed_tokens + + scheduler_output.num_scheduled_tokens[request_id] + ) + if not self.save_decode_cache: + total_tokens = min( + total_tokens, cached_request.num_prompt_tokens + ) + store_plans = self._store_plans( + cached_request, + cached_tracker, + total_tokens, + block_ids_override=cached_tracker.block_ids, + ) + meta.store_plans.extend(store_plans) + if store_plans: + meta.store_requests[request_id] = store_plans + + # Async loads are admitted without scheduling model tokens, so those + # requests are absent from both scheduled request collections above. + # Forward their plans explicitly so workers can start the transfer and + # eventually move the request out of WAITING_FOR_REMOTE_KVS. + for request_id, load_plans in list(self._pending_loads.items()): + self._pending_loads.pop(request_id, None) + self._load_specs.pop(request_id, None) + if load_plans: + meta.load_requests[request_id] = load_plans + + for request_id in scheduler_output.preempted_req_ids or set(): + meta.preempted_request_ids.add(request_id) + self._load_specs.pop(request_id, None) + if preempted_tracker := self._request_trackers.get(request_id): + preempted_tracker.reset() + + event_plans = tuple(meta.store_plans) + if event_plans: + event = self._store_event_counter + self._store_event_counter += 1 + meta.store_event = event + self._store_events[event] = _StoreEvent(event_plans) + if self._gpu_block_pool is not None: + self._gpu_block_pool.touch( + [ + self._gpu_block_pool.blocks[block_id] + for block_id in dict.fromkeys( + plan.block_id for plan in event_plans + ) + ] + ) + return meta + + def _tracker_for_request(self, request_id: str) -> RequestTracker: + return self._request_trackers.setdefault(request_id, RequestTracker()) + + def _load_plans_for_external_tokens( + self, + request: Any, + tracker: RequestTracker, + block_groups: tuple[list[int], ...], + num_external_tokens: int, + ) -> list[BlockTransferPlan]: + """Plan the logical objects selected by the external hit window.""" + group_ids = self.kv_cache_config.prefix_cacheable_group_ids + spec = tracker.load_spec + local_tokens = spec.local_tokens if spec is not None else 0 + end_tokens = local_tokens + num_external_tokens + plans: list[BlockTransferPlan] = [] + for group_index, group_id in enumerate(group_ids): + group_block_size = self.group_block_sizes[group_id] + start_block = local_tokens // group_block_size + num_full_blocks = end_tokens // group_block_size + for block_index in range(start_block, num_full_blocks): + token_end = (block_index + 1) * group_block_size + block_hash = self._object_hash_at_token_end( + request.block_hashes, token_end + ) + plans.append( + BlockTransferPlan( + key=self.codec.key(block_hash, group_id=group_id), + block_id=block_groups[group_index][block_index], + request_id=request.request_id, + group_id=group_id, + ) + ) + return plans + + def _object_hash_at_token_end( + self, hashes: Sequence[bytes], token_end: int + ) -> bytes: + hash_index = token_end // self.hash_block_size - 1 + if hash_index < 0 or hash_index >= len(hashes): + raise ValueError( + f"request has {len(hashes)} hashes, needs index {hash_index}" + ) + return hashes[hash_index] + + def _store_plans( + self, + request: Any, + tracker: RequestTracker, + token_count: int, + block_ids_override: tuple[list[int], ...] | None = None, + ) -> list[BlockTransferPlan]: + """Describe full blocks produced by a scheduled prefill.""" + request_id = getattr(request, "request_id", None) or request.req_id + group_ids = self.kv_cache_config.prefix_cacheable_group_ids + # Scheduled speculative tokens are not yet verified and have no block + # hash; a block is storable only once its hash exists. + token_count = min(token_count, len(request.block_hashes) * self.hash_block_size) + previous_saved_tokens = tracker.saved_tokens + # Hash-aligned rather than block-aligned: a group with smaller blocks + # than the largest must not wait for that group's block to fill. + save_to = tracker.mark_saved(token_count, self.hash_block_size) + plans: list[BlockTransferPlan] = [] + block_groups = ( + block_ids_override + if block_ids_override is not None + else tuple(request.block_ids[group_id] for group_id in group_ids) + ) + block_ranges = { + group_id: range( + previous_saved_tokens // self.group_block_sizes[group_id], + min(save_to // self.group_block_sizes[group_id], len(block_ids)), + ) + for group_id, block_ids in zip(group_ids, block_groups) + } + for group_id, block_ids in zip(group_ids, block_groups): + if group_id not in block_ranges: + continue + group_block_size = self.group_block_sizes[group_id] + for index in block_ranges[group_id]: + if block_ids[index] == NULL_BLOCK_ID: + continue + token_end = (index + 1) * group_block_size + block_hash = self._object_hash_at_token_end( + request.block_hashes, token_end + ) + plans.append( + BlockTransferPlan( + key=self.codec.key(block_hash, group_id=group_id), + block_id=block_ids[index], + request_id=request_id, + group_id=group_id, + ) + ) + return plans + + def request_finished( + self, request: Request, block_ids: tuple[list[int], ...] + ) -> tuple[bool, dict[str, Any] | None]: + self._pending_loads.pop(request.request_id, None) + self._load_specs.pop(request.request_id, None) + self._requests.pop(request.request_id, None) + return False, None + + def update_connector_output(self, output: KVConnectorOutput) -> None: + metadata = output.kv_connector_worker_meta + if not isinstance(metadata, UMBPConnectorWorkerMetadata): + return + pool = self._gpu_block_pool + for event_id, result in metadata.store_events.items(): + event = self._store_events.get(event_id) + if event is None: + continue + event.result.merge(result) + if event.result.completed_workers < self._num_workers: + continue + del self._store_events[event_id] + if pool is not None: + block_ids = dict.fromkeys(plan.block_id for plan in event.plans) + pool.free_blocks( + pool.blocks[block_id] for block_id in reversed(block_ids) + ) + + def has_pending_push_work(self) -> bool: + return bool(self._store_events) + + def reset_store(self) -> bool: + if self._store_events: + return False + self._load_specs.clear() + self._pending_loads.clear() + return self.runtime.clear() + + def close(self) -> None: + self.runtime.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py new file mode 100644 index 0000000000..34101f294c --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared worker-side UMBP connector behavior.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +import torch + +from vllm.forward_context import ForwardContext +from vllm.logger import init_logger + +from .data import ( + BlockIdentityCodec, + BlockTransferPlan, + KVLayoutPlanner, + StoreEventResult, + TransferJobState, + TransferJobStatus, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, +) +from .runtime import UMBPWorkerHandle + +logger = init_logger(__name__) + + +@dataclass +class _StoreBatch: + event_id: int + jobs: dict[str, TransferJobState] + + +class UMBPStoreConnectorWorker: + """Translate shared metadata into runtime load/store jobs.""" + + def __init__( + self, + runtime: UMBPWorkerHandle, + layout: KVLayoutPlanner | None = None, + *, + codec: BlockIdentityCodec | None = None, + ) -> None: + self.runtime = runtime + self.layout = layout + self.codec = codec + self._load_jobs: dict[str, dict[str | None, TransferJobState]] = {} + self._store_jobs: dict[str, TransferJobState] = {} + self._pending_stores: list[_StoreBatch] = [] + self._worker_meta = UMBPConnectorWorkerMetadata() + self._finished_recving: set[str] = set() + self._load_error_block_ids: set[int] = set() + self._report_load_completions = False + self._active_store_event = -1 + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + if self.layout is not None and self.layout.regions: + self.layout.register_kv_caches(kv_caches) + self.runtime.register_buffers(kv_caches) + + def _materialize_plans( + self, plans: list[BlockTransferPlan] + ) -> list[BlockTransferPlan]: + plans = [self._localize_plan(plan) for plan in plans] + if self.layout is None or not self.layout.regions: + return plans + materialized: list[BlockTransferPlan] = [] + for plan in plans: + if plan.ranges: + materialized.append(plan) + continue + materialized.append(self.layout.materialize(plan)) + return materialized + + def _localize_plan(self, plan: BlockTransferPlan) -> BlockTransferPlan: + if self.codec is None or plan.group_id is None: + return plan + key = self._localize_key(plan.key, plan.group_id) + if key == plan.key: + return plan + return replace(plan, key=key) + + def _localize_key(self, key: str, group_id: int) -> str: + if self.codec is None: + return key + try: + block_hash = bytes.fromhex(key.rsplit(":", 1)[1]) + except (IndexError, ValueError): + return key + return self.codec.key(block_hash, group_id) + + def start_load_kv( + self, forward_context: ForwardContext, metadata: UMBPConnectorMetadata + ) -> None: + del forward_context + self._report_load_completions = metadata.async_load + for request_id, plans in metadata.load_requests.items(): + if plans: + self._submit_loads(request_id, plans) + + def _submit_loads(self, request_id: str, plans: list[BlockTransferPlan]) -> None: + materialized = self._materialize_plans(plans) + logger.debug( + "UMBP load submitted request=%s plans=%d ranges=%d bytes=%d", + request_id, + len(materialized), + sum(len(plan.ranges) for plan in materialized), + sum(item.length for plan in materialized for item in plan.ranges), + ) + self._load_jobs[request_id] = {None: self.runtime.load(materialized)} + + def _finish_job( + self, request_id: str, result: TransferJobState, *, is_load: bool + ) -> bool: + operation = "load" if is_load else "store" + succeeded = result.status == TransferJobStatus.COMPLETED + if not succeeded: + logger.debug( + "UMBP %s failed request=%s status=%s error=%s", + operation, + request_id, + result.status.value, + result.error, + ) + if is_load: + self._load_error_block_ids.update(result.failed_block_ids) + elif succeeded: + self.runtime.publish(result) + return succeeded + + def wait_for_layer_load(self, layer_name: str) -> None: + if self._report_load_completions: + # Asynchronous loads fill blocks of requests that do not run until + # get_finished reports them, so no forward pass reads those blocks. + return + self._drain_load_jobs(wait=True) + + def wait_for_save(self) -> None: + if self._store_jobs or self._active_store_event >= 0: + self._pending_stores.append( + _StoreBatch(self._active_store_event, self._store_jobs) + ) + self._store_jobs = {} + self._active_store_event = -1 + self._drain_store_jobs(wait=False) + + def _drain_store_jobs(self, *, wait: bool) -> None: + for batch in list(self._pending_stores): + for request_id, job in list(batch.jobs.items()): + result = self.runtime.wait(job) if wait else self.runtime.poll(job) + if result is None: + continue + # A partly failed job is not published, including its stored keys. + self._finish_job(request_id, result, is_load=False) + del batch.jobs[request_id] + if not batch.jobs: + if batch.event_id >= 0: + self._worker_meta.store_events[batch.event_id] = StoreEventResult( + completed_workers=1 + ) + self._pending_stores.remove(batch) + + def _submit_store_plans( + self, plans: list[BlockTransferPlan], request_id: str + ) -> None: + if not plans: + return + materialized = self._materialize_plans(plans) + self._store_jobs[request_id] = self.runtime.store(materialized) + + def enqueue_stores(self, metadata: UMBPConnectorMetadata) -> None: + self._active_store_event = metadata.store_event + store_requests = { + request_id: list(plans) + for request_id, plans in metadata.store_requests.items() + } + grouped = {plan for plans in store_requests.values() for plan in plans} + for plan in metadata.store_plans: + if plan not in grouped: + store_requests.setdefault( + plan.request_id or "__umbp_batch__", [] + ).append(plan) + grouped.add(plan) + for request_id, plans in store_requests.items(): + self._submit_store_plans(plans, request_id) + + def handle_preemptions(self, metadata: UMBPConnectorMetadata) -> None: + """Finish request-local jobs before vLLM reuses their GPU blocks.""" + preempted = metadata.preempted_request_ids + if not preempted: + return + for request_id in preempted: + for job in self._load_jobs.pop(request_id, {}).values(): + self.runtime.wait(job) + for jobs in [self._store_jobs, *(batch.jobs for batch in self._pending_stores)]: + for request_id, job in jobs.items(): + if request_id in preempted or any( + plan.request_id in preempted for plan in job.plans + ): + jobs[request_id] = self.runtime.wait(job) + + def get_finished( + self, finished_req_ids: set[str] + ) -> tuple[set[str] | None, set[str] | None]: + del finished_req_ids + self._drain_load_jobs(wait=False) + self._drain_store_jobs(wait=False) + recving = set(self._finished_recving) + self._finished_recving.difference_update(recving) + return None, recving or None + + def _drain_load_jobs(self, *, wait: bool) -> None: + for request_id, jobs in list(self._load_jobs.items()): + for layer, job in list(jobs.items()): + result = self.runtime.wait(job) if wait else self.runtime.poll(job) + if result is None: + continue + del jobs[layer] + self._finish_job(request_id, result, is_load=True) + if not jobs: + del self._load_jobs[request_id] + if self._report_load_completions: + self._finished_recving.add(request_id) + + def build_connector_worker_meta(self) -> UMBPConnectorWorkerMetadata: + result = self._worker_meta + self._worker_meta = UMBPConnectorWorkerMetadata() + return result + + def get_block_ids_with_load_errors(self) -> set[int]: + failed = self._load_error_block_ids + self._load_error_block_ids = set() + return failed + + def close(self) -> None: + self.wait_for_save() + self._drain_store_jobs(wait=True) + self._drain_load_jobs(wait=True) + self.runtime.close()