[UMBP][1/10] Add shared connector, runtime contracts, and request lifecycle (#61789)
Signed-off-by: CharlesWu <Charles.Wu@amd.com> Co-authored-by: Yichao Zhu <Yichao.Zhu@amd.com>
This commit is contained in:
CharlesWu
@@ -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",
|
||||
]
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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."""
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user