vllm.distributed.kv_transfer.kv_connector.v1.umbp.data

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.

Classes:

• BlockIdentityCodec –

  Build stable rank-local keys from vLLM block hashes.

• BlockLoadBatch –

  Whole-block loads with request metadata shared across the batch.

• BlockStoreBatch –

  Columnar whole-block stores with request metadata shared by the batch.

• BlockTransferPlan –

  One complete logical block transfer.

• KVLayoutDescriptor –

  Backend-neutral description of one rank-local KV object layout.

• KVLayoutPlanner –

  Create a deterministic all-layer object layout from KVCacheConfig.

• KVRange –

  A source/destination range in a runtime transfer.

• KVRegion –

  One deterministic region in a rank-local KV object.

• LoadSpec –

  Scheduler decision for one externally loaded prefix.

• RankTopology –

  Model-parallel ranks that participate in one KV object.

• RequestTracker –

  Core request state shared by every runtime mode.

• StoreEventResult –

  Number of workers that finished one store event.

• TransferJobState –

  Per-key completion state; one failed key does not fail its siblings.

• UMBPConnectorMetadata –

  Scheduler-to-worker metadata for one engine step.

• UMBPConnectorWorkerMetadata –

  Worker completion metadata aggregated across ranks.

• UMBPHandshakeMetadata –

  Namespace a worker derived, reported to the scheduler at startup.

• UMBPNamespace –

  Compatibility namespace for objects stored outside the engine.

Functions:

• tp_replicated_group_ids –

  Groups whose blocks are byte-identical on every TP rank.

BlockIdentityCodec dataclass

Build stable rank-local keys from vLLM block hashes.

Groups in replicated_groups get one TP-agnostic key (tp*) that every TP rank loads and only TP rank 0 stores.

Methods:

• keys_per_object –

  Number of rank keys that together hold one logical block.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass(frozen=True)
class BlockIdentityCodec:
    """Build stable rank-local keys from vLLM block hashes.

    Groups in `replicated_groups` get one TP-agnostic key (`tp*`) that
    every TP rank loads and only TP rank 0 stores.
    """

    namespace: UMBPNamespace
    tp_rank: int = 0
    pp_rank: int = 0
    pcp_rank: int = 0
    dcp_rank: int = 0
    replicated_groups: frozenset[int] = frozenset()

    def stores_group(self, group_id: int) -> bool:
        return self.tp_rank == 0 or group_id not in self.replicated_groups

    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
        tp = "*" if group_id in self.replicated_groups else tp_rank
        return (
            f"umbp:vllm:v1:{self.namespace.value}:"
            f"tp{tp}: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
            if group_id not in self.replicated_groups or namespace[0] == 0
        )

    def keys_per_object(self, topology: RankTopology, group_id: int) -> int:
        """Number of rank keys that together hold one logical block."""
        if group_id in self.replicated_groups:
            return topology.rank_count // topology.tp_size
        return topology.rank_count

keys_per_object(topology, group_id)

Number of rank keys that together hold one logical block.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

def keys_per_object(self, topology: RankTopology, group_id: int) -> int:
    """Number of rank keys that together hold one logical block."""
    if group_id in self.replicated_groups:
        return topology.rank_count // topology.tp_size
    return topology.rank_count

BlockLoadBatch dataclass

Bases: Sequence[BlockTransferPlan]

Whole-block loads with request metadata shared across the batch.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass(frozen=True)
class BlockLoadBatch(Sequence[BlockTransferPlan]):
    """Whole-block loads with request metadata shared across the batch."""

    keys: list[str]
    block_ids: list[int]
    group_ids: list[int]
    request_id: str

    def __post_init__(self) -> None:
        if len(self.keys) != len(self.block_ids) or len(self.keys) != len(
            self.group_ids
        ):
            raise ValueError("load batch columns must have equal lengths")

    def __len__(self) -> int:
        return len(self.keys)

    @overload
    def __getitem__(self, index: int) -> BlockTransferPlan: ...

    @overload
    def __getitem__(self, index: slice) -> list[BlockTransferPlan]: ...

    def __getitem__(
        self, index: int | slice
    ) -> BlockTransferPlan | list[BlockTransferPlan]:
        if isinstance(index, slice):
            return [self[i] for i in range(*index.indices(len(self)))]
        return BlockTransferPlan(
            key=self.keys[index],
            block_id=self.block_ids[index],
            group_id=self.group_ids[index],
            request_id=self.request_id,
        )

BlockStoreBatch dataclass

Bases: Sequence[BlockTransferPlan]

Columnar whole-block stores with request metadata shared by the batch.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass(frozen=True)
class BlockStoreBatch(Sequence[BlockTransferPlan]):
    """Columnar whole-block stores with request metadata shared by the batch."""

    keys: list[str]
    block_ids: list[int]
    group_ids: list[int]
    block_hashes: list[bytes]
    parent_block_hashes: list[bytes | None]
    token_ids: list[tuple[int, ...]]
    block_sizes: list[int]
    request_id: str

    def __post_init__(self) -> None:
        length = len(self.keys)
        columns = (
            self.block_ids,
            self.group_ids,
            self.block_hashes,
            self.parent_block_hashes,
            self.token_ids,
            self.block_sizes,
        )
        if any(len(column) != length for column in columns):
            raise ValueError("store batch columns must have equal lengths")

    def __len__(self) -> int:
        return len(self.keys)

    @overload
    def __getitem__(self, index: int) -> BlockTransferPlan: ...

    @overload
    def __getitem__(self, index: slice) -> list[BlockTransferPlan]: ...

    def __getitem__(
        self, index: int | slice
    ) -> BlockTransferPlan | list[BlockTransferPlan]:
        if isinstance(index, slice):
            return [self[i] for i in range(*index.indices(len(self)))]
        return BlockTransferPlan(
            key=self.keys[index],
            block_id=self.block_ids[index],
            request_id=self.request_id,
            group_id=self.group_ids[index],
            block_hash=self.block_hashes[index],
            parent_block_hash=self.parent_block_hashes[index],
            token_ids=self.token_ids[index],
            block_size=self.block_sizes[index],
        )

    def localized(self, keys: list[str]) -> BlockStoreBatch:
        return replace(self, keys=keys)

BlockTransferPlan dataclass

One complete logical block transfer.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@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
    block_hash: bytes | None = None
    parent_block_hash: bytes | None = None
    token_ids: tuple[int, ...] = ()
    block_size: int = 0
    medium: str = "CPU"

KVLayoutDescriptor dataclass

Backend-neutral description of one rank-local KV object layout.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@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)

KVLayoutPlanner

Create a deterministic all-layer object layout from KVCacheConfig.

Methods:

• plan_for_block –

  Lower one GPU block into a deterministic scatter/gather plan.

• register_kv_caches –

  Capture stable base addresses for the planned KV regions.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

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:
        layer_specs = {
            layer: (group_id, spec)
            for group_id, group in enumerate(config.kv_cache_groups)
            if group.enable_kv_transfer
            for layer, spec in _layer_specs(group).items()
        }
        layout = config.kv_cache_layout
        regions: list[KVRegion] = []
        offset = 0
        for tensor in config.kv_cache_tensors:
            for layer_name in tensor.layers:
                if layer_name not in layer_specs:
                    continue
                group_id, spec = layer_specs[layer_name]
                block_bytes = spec.page_size_bytes
                segments = (
                    ()
                    if layout is None
                    else _page_segments(
                        spec, tensor, config.num_blocks, KVCacheLayout[layout]
                    )
                )
                if not segments and (
                    block_bytes > tensor.block_stride
                    or (len(tensor.layers) > 1 and tensor.layer_stride < block_bytes)
                ):
                    raise ValueError(
                        f"UMBP cannot copy {layer_name!r}: its KV block is split "
                        "into strided runs and KVCacheConfig.kv_cache_layout is unset"
                    )
                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,
                        segments=segments,
                    )
                )
                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 region.segments and block_stride != region.block_stride:
                    raise ValueError(
                        f"KV block stride mismatch for {layer_name!r}: "
                        f"config={region.block_stride}, tensor={block_stride}"
                    )
                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.segments:
                start = object_offset if group_id is not None else region.object_offset
                for segment_offset, length in region.segments:
                    ranges.append(
                        KVRange(
                            layer_name=region.layer_name,
                            group_id=region.group_id,
                            block_id=block_id,
                            base_address=base_address
                            + block_id * region.block_stride
                            + segment_offset,
                            stride=region.block_stride,
                            length=length,
                            object_offset=start,
                        )
                    )
                    start += length
                object_offset += region.block_bytes
                continue
            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)

plan_for_block(key, block_id, base_addresses, *, group_id=None, request_id=None)

Lower one GPU block into a deterministic scatter/gather plan.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

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,
    )

register_kv_caches(kv_caches)

Capture stable base addresses for the planned KV regions.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

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 region.segments and block_stride != region.block_stride:
                raise ValueError(
                    f"KV block stride mismatch for {layer_name!r}: "
                    f"config={region.block_stride}, tensor={block_stride}"
                )
            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

KVRange

Bases: NamedTuple

A source/destination range in a runtime transfer.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

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

KVRegion dataclass

One deterministic region in a rank-local KV object.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@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
    segments: tuple[tuple[int, int], ...] = ()

    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.segments:
            if sum(length for _, length in self.segments) != self.block_bytes:
                raise ValueError(
                    f"KV region segments of {self.layer_name!r} do not cover "
                    f"its {self.block_bytes} bytes"
                )
        elif 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}"
            )

LoadSpec dataclass

Scheduler decision for one externally loaded prefix.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass
class LoadSpec:
    """Scheduler decision for one externally loaded prefix."""

    local_tokens: int
    external_tokens: int
    # Group block indices start at the beginning of the request.
    block_hashes_by_group: tuple[tuple[bytes | None, ...], ...] = ()

    @property
    def num_tokens_to_load(self) -> int:
        return max(self.external_tokens - self.local_tokens, 0)

RankTopology dataclass

Model-parallel ranks that participate in one KV object.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@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)
        )

RequestTracker dataclass

Core request state shared by every runtime mode.

Methods:

• mark_saved –

  Advance the save watermark only to a complete block boundary.

• replace_blocks –

  Replace stale pre-preemption tables with resumed full tables.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@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

mark_saved(token_count, block_size)

Advance the save watermark only to a complete block boundary.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

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

replace_blocks(block_ids)

Replace stale pre-preemption tables with resumed full tables.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

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

StoreEventResult dataclass

Number of workers that finished one store event.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass
class StoreEventResult:
    """Number of workers that finished one store event."""

    completed_workers: int = 0
    failed_objects: set[tuple[int, bytes]] = field(default_factory=set)

    def merge(self, other: StoreEventResult) -> None:
        self.completed_workers += other.completed_workers
        self.failed_objects.update(other.failed_objects)

TransferJobState dataclass

Per-key completion state; one failed key does not fail its siblings.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass
class TransferJobState:
    """Per-key completion state; one failed key does not fail its siblings."""

    plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch | BlockStoreBatch
    status: TransferJobStatus = TransferJobStatus.PENDING
    completed_keys: set[str] = field(default_factory=set)
    failed_keys: set[str] = field(default_factory=set)
    error: str | None = None
    plan_bytes: tuple[int, ...] | None = None

    @property
    def keys(self) -> Sequence[str]:
        if isinstance(self.plans, (BlockLoadBatch, BlockStoreBatch)):
            return self.plans.keys
        return tuple(plan.key for plan in self.plans)

    @property
    def completed_bytes(self) -> int:
        if self.plan_bytes is not None:
            return sum(
                size
                for key, size in zip(self.keys, self.plan_bytes, strict=True)
                if key in self.completed_keys
            )
        return sum(
            item.length
            for plan in self.plans
            if plan.key in self.completed_keys
            for item in plan.ranges
        )

    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()

    def cancel(self, error: str = "cancelled") -> None:
        self.error = error
        self.status = TransferJobStatus.CANCELLED

    @property
    def failed_block_ids(self) -> set[int]:
        if isinstance(self.plans, (BlockLoadBatch, BlockStoreBatch)):
            return {
                block_id
                for key, block_id in zip(
                    self.plans.keys, self.plans.block_ids, strict=True
                )
                if key in self.failed_keys
            }
        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
            )

UMBPConnectorMetadata dataclass

Bases: KVConnectorMetadata

Scheduler-to-worker metadata for one engine step.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@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] | BlockLoadBatch] = field(
        default_factory=dict
    )
    store_requests: dict[str, list[BlockTransferPlan] | BlockStoreBatch] = field(
        default_factory=dict
    )
    preempted_request_ids: set[str] = field(default_factory=set)
    store_event: int = -1

    def iter_store_plans(self):
        yield from self.store_plans
        for plans in self.store_requests.values():
            yield from plans

    def store_plan_groups(self):
        if self.store_plans:
            yield self.store_plans
        yield from self.store_requests.values()

UMBPConnectorWorkerMetadata dataclass

Bases: KVConnectorWorkerMetadata

Worker completion metadata aggregated across ranks.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@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

UMBPHandshakeMetadata dataclass

Bases: KVConnectorHandshakeMetadata

Namespace a worker derived, reported to the scheduler at startup.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass(frozen=True)
class UMBPHandshakeMetadata(KVConnectorHandshakeMetadata):
    """Namespace a worker derived, reported to the scheduler at startup."""

    namespace: str
    replicated_groups: tuple[int, ...] = ()

UMBPNamespace dataclass

Compatibility namespace for objects stored outside the engine.

Methods:

• configured –

  Return the explicit key_namespace, or None when it is derived.

• from_vllm_config –

  Derive the namespace from a worker's per-layer KV cache config.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@dataclass(frozen=True)
class UMBPNamespace:
    """Compatibility namespace for objects stored outside the engine."""

    value: str

    @classmethod
    def configured(cls, vllm_config: Any) -> UMBPNamespace | None:
        """Return the explicit `key_namespace`, or None when it is derived."""
        extra = vllm_config.kv_transfer_config.kv_connector_extra_config
        configured = extra.get("key_namespace", "auto")
        if configured == "auto":
            return None
        if not isinstance(configured, str) or not configured:
            raise ValueError("key_namespace must be a non-empty string or 'auto'")
        return cls(configured)

    @classmethod
    def from_vllm_config(
        cls,
        vllm_config: Any,
        kv_cache_config: KVCacheConfig,
    ) -> UMBPNamespace:
        """Derive the namespace from a worker's per-layer KV cache config.

        The scheduler's config flattens UniformTypeKVCacheSpecs groups, so the
        scheduler adopts the namespace its workers report instead.
        """
        if (configured := cls.configured(vllm_config)) is not None:
            return 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),
            "model_config": _model_fingerprint(model_config),
            "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,
                    "layer_specs": {
                        name: _spec_fingerprint(spec)
                        for name, spec in _layer_specs(group).items()
                    },
                }
                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": 2,
        }
        if speculative is not None:
            # Added only when set, so namespaces without a drafter are unchanged.
            fields["speculative"] = speculative
        if _engine_may_load_foreign_kv(vllm_config):
            # Bytes another connector loaded can be stored here too. Keep them
            # out of the namespace that UMBP-only engines restore as canonical
            # KV. Added only when set, so UMBP-only namespaces are unchanged.
            fields["foreign_kv"] = True
        payload = json.dumps(
            fields, sort_keys=True, separators=(",", ":"), default=str
        ).encode()
        return cls(hashlib.sha256(payload).hexdigest()[:32])

configured(vllm_config) classmethod

Return the explicit key_namespace, or None when it is derived.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@classmethod
def configured(cls, vllm_config: Any) -> UMBPNamespace | None:
    """Return the explicit `key_namespace`, or None when it is derived."""
    extra = vllm_config.kv_transfer_config.kv_connector_extra_config
    configured = extra.get("key_namespace", "auto")
    if configured == "auto":
        return None
    if not isinstance(configured, str) or not configured:
        raise ValueError("key_namespace must be a non-empty string or 'auto'")
    return cls(configured)

from_vllm_config(vllm_config, kv_cache_config) classmethod

Derive the namespace from a worker's per-layer KV cache config.

The scheduler's config flattens UniformTypeKVCacheSpecs groups, so the scheduler adopts the namespace its workers report instead.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

@classmethod
def from_vllm_config(
    cls,
    vllm_config: Any,
    kv_cache_config: KVCacheConfig,
) -> UMBPNamespace:
    """Derive the namespace from a worker's per-layer KV cache config.

    The scheduler's config flattens UniformTypeKVCacheSpecs groups, so the
    scheduler adopts the namespace its workers report instead.
    """
    if (configured := cls.configured(vllm_config)) is not None:
        return 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),
        "model_config": _model_fingerprint(model_config),
        "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,
                "layer_specs": {
                    name: _spec_fingerprint(spec)
                    for name, spec in _layer_specs(group).items()
                },
            }
            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": 2,
    }
    if speculative is not None:
        # Added only when set, so namespaces without a drafter are unchanged.
        fields["speculative"] = speculative
    if _engine_may_load_foreign_kv(vllm_config):
        # Bytes another connector loaded can be stored here too. Keep them
        # out of the namespace that UMBP-only engines restore as canonical
        # KV. Added only when set, so UMBP-only namespaces are unchanged.
        fields["foreign_kv"] = True
    payload = json.dumps(
        fields, sort_keys=True, separators=(",", ":"), default=str
    ).encode()
    return cls(hashlib.sha256(payload).hexdigest()[:32])

_engine_may_load_foreign_kv(vllm_config)

Whether any connector of this engine may load KV that vLLM's own cache writer did not produce, e.g. a MultiConnector sibling of UMBP.

MultiConnector hands each child a config copy holding only that child's kv_transfer_config, so read the engine's current config, which workers set while they create their connector.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

def _engine_may_load_foreign_kv(vllm_config: Any) -> bool:
    """Whether any connector of this engine may load KV that vLLM's own cache
    writer did not produce, e.g. a MultiConnector sibling of UMBP.

    MultiConnector hands each child a config copy holding only that child's
    kv_transfer_config, so read the engine's current config, which workers set
    while they create their connector.
    """
    from vllm.config import get_current_vllm_config_or_none
    from vllm.distributed.kv_transfer.kv_connector.utils import (
        kv_connector_may_load_foreign_kv,
    )

    engine_config = get_current_vllm_config_or_none()
    if engine_config is None or engine_config.kv_transfer_config is None:
        engine_config = vllm_config
    return kv_connector_may_load_foreign_kv(engine_config)

_layer_specs(group)

Map each layer of a cache group to the spec describing its own pages.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

def _layer_specs(group: KVCacheGroupSpec) -> dict[str, KVCacheSpec]:
    """Map each layer of a cache group to the spec describing its own pages."""
    spec = group.kv_cache_spec
    if isinstance(spec, UniformTypeKVCacheSpecs):
        return {name: spec.kv_cache_specs[name] for name in group.layer_names}
    return dict.fromkeys(group.layer_names, spec)

_model_fingerprint(model_config)

Effective model settings that change the KV a checkpoint produces.

hf_config already has hf_overrides applied, so overridden RoPE or attention settings separate engines that share the same model path.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

def _model_fingerprint(model_config: Any) -> dict[str, Any]:
    """Effective model settings that change the KV a checkpoint produces.

    `hf_config` already has `hf_overrides` applied, so overridden RoPE or
    attention settings separate engines that share the same model path.
    """
    hf_config = getattr(model_config, "hf_config", None)
    to_dict = getattr(hf_config, "to_dict", None)
    config = dict(to_dict()) if callable(to_dict) else {}
    # Neither changes the KV, and either would split engines sharing a pool.
    config.pop("_name_or_path", None)
    config.pop("transformers_version", None)
    return {
        "dtype": str(getattr(model_config, "dtype", None)),
        "quantization": getattr(model_config, "quantization", None),
        "hf_config": config,
    }

_page_segments(spec, tensor, num_blocks, layout)

Byte (offset, length) runs of a layer's page from its block start.

Empty when the page is one contiguous run. LHBNC and BHLNC store each head's tokens as a separate run.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

def _page_segments(
    spec: KVCacheSpec,
    tensor: KVCacheTensor,
    num_blocks: int,
    layout: KVCacheLayout,
) -> tuple[tuple[int, int], ...]:
    """Byte `(offset, length)` runs of a layer's page from its block start.

    Empty when the page is one contiguous run. LHBNC and BHLNC store each
    head's tokens as a separate run.
    """
    if not spec.has_layer_views:
        return ()
    _, heads, states, content = compute_layer_kv_cache_shape_bytes(spec, num_blocks)
    _, _, head_stride, state_stride, _ = compute_layout_strides(
        spec,
        num_blocks,
        len(tensor.layers),
        layout,
        fixed_strides=(tensor.layer_stride, tensor.block_stride, None, None, None),
    )
    dims = sorted(
        (stride, size)
        for stride, size in ((head_stride, heads), (state_stride, states))
        if size > 1
    )
    run = content
    while dims and dims[0][0] == run:
        run *= dims.pop(0)[1]
    if not dims:
        return ()
    return tuple(
        (sum(i * stride for i, (stride, _) in zip(index, dims)), run)
        for index in itertools.product(*(range(size) for _, size in dims))
    )

_spec_fingerprint(spec)

Every field that shapes a layer's stored bytes, in a stable form.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

def _spec_fingerprint(spec: KVCacheSpec) -> dict[str, str]:
    """Every field that shapes a layer's stored bytes, in a stable form."""
    return {
        "kind": type(spec).__name__,
        **{item.name: str(getattr(spec, item.name)) for item in dataclass_fields(spec)},
    }

tp_replicated_group_ids(config, topology)

Groups whose blocks are byte-identical on every TP rank.

Every layer must declare a single TP shard (max_tp_shards=1, e.g. MLA latents and sparse indexers); context parallelism shards them, so it opts out.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py

def tp_replicated_group_ids(
    config: KVCacheConfig, topology: RankTopology
) -> frozenset[int]:
    """Groups whose blocks are byte-identical on every TP rank.

    Every layer must declare a single TP shard (`max_tp_shards=1`, e.g. MLA
    latents and sparse indexers); context parallelism shards them, so it opts
    out.
    """
    if topology.tp_size == 1 or topology.pcp_size > 1 or topology.dcp_size > 1:
        return frozenset()
    return frozenset(
        group_id
        for group_id, group in enumerate(config.kv_cache_groups)
        if group.enable_kv_transfer
        and group.layer_names
        and all(
            isinstance(spec, AttentionSpec) and spec.max_tp_shards == 1
            for spec in _layer_specs(group).values()
        )
    )