vllm.distributed.kv_transfer.kv_connector.v1.umbp

Modules:

• connector –

  vLLM-facing shared UMBP connector.

• data –

  Shared data types for the UMBP KV connector.

• runtime –
• scheduler –

  Shared scheduler-side UMBP connector behavior.

• worker –

  Shared worker-side UMBP connector behavior.

Classes:

• BlockIdentityCodec –

  Build stable rank-local keys from vLLM block hashes.

• 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.

• 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.

• UMBPNamespace –

  Compatibility namespace for objects stored outside the engine.

• UMBPRuntimeConfig –

  Shared connector options passed to the selected runtime adapter.

• UMBPRuntimeFactory –

  Registry used by UMBP adapters without importing MORI in vLLM.

• UMBPStoreConnector –

  vLLM connector lifecycle shared by every UMBP runtime mode.

BlockIdentityCodec dataclass

Build stable rank-local keys from vLLM block hashes.

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."""

    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
        )

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

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

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

    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}"
            )

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

    @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

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

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]] = 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

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

UMBPNamespace dataclass

Compatibility namespace for objects stored outside the engine.

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

UMBPRuntimeConfig dataclass

Shared connector options passed to the selected runtime adapter.

Methods:

• resolve_for_rank_count –

  Pass topology to the adapter without interpreting its capacity options.

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

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

resolve_for_rank_count(rank_count)

Pass topology to the adapter without interpreting its capacity options.

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

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)

UMBPRuntimeFactory

Registry used by UMBP adapters without importing MORI in vLLM.

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

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)

UMBPStoreConnector

Bases: KVConnectorBase_V1

vLLM connector lifecycle shared by every UMBP runtime mode.

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

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