Skip to content

vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime

Modules:

  • base –

    Backend-neutral runtime contract for UMBP.

  • factory –

    Runtime configuration and adapter factory for UMBP.

Classes:

IUMBPRuntime

Bases: Protocol

Methods:

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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."""

create_scheduler_handle(namespace, topology, layout)

Create a scheduler-side handle.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
55
56
57
58
59
60
61
def create_scheduler_handle(
    self,
    namespace: str,
    topology: RankTopology,
    layout: KVLayoutDescriptor,
) -> UMBPSchedulerHandle:
    """Create a scheduler-side handle."""

create_worker_handle(namespace, topology, layout)

Create a worker-side handle.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
63
64
65
66
67
68
69
def create_worker_handle(
    self,
    namespace: str,
    topology: RankTopology,
    layout: KVLayoutDescriptor,
) -> UMBPWorkerHandle:
    """Create a worker-side handle."""

UMBPRuntimeConfig dataclass

Shared connector options passed to the selected runtime adapter.

Methods:

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@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 
45
46
47
48
49
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 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
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)

UMBPSchedulerHandle

Bases: Protocol

Methods:

  • clear –

    Clear all objects visible to this scheduler handle.

  • close –

    Release scheduler-side resources.

  • lookup –

    Return authoritative existence for each key.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
20
21
22
23
24
25
26
27
28
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."""

clear()

Clear all objects visible to this scheduler handle.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
24
25
def clear(self) -> bool:
    """Clear all objects visible to this scheduler handle."""

close()

Release scheduler-side resources.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
27
28
def close(self) -> None:
    """Release scheduler-side resources."""

lookup(keys)

Return authoritative existence for each key.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
21
22
def lookup(self, keys: Sequence[str]) -> Sequence[bool]:
    """Return authoritative existence for each key."""

UMBPWorkerHandle

Bases: Protocol

Methods:

  • close –

    Release worker-side resources.

  • load –

    Start loading the requested plans.

  • poll –

    Return a finished job without blocking, or None if still pending.

  • publish –

    Accept a completed store; a runtime may defer visibility until then.

  • register_buffers –

    Register each unique KV allocation with the runtime.

  • store –

    Start storing the requested plans.

  • wait –

    Block until the job is final and its buffers are safe to reuse.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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."""

close()

Release worker-side resources.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
50
51
def close(self) -> None:
    """Release worker-side resources."""

load(plans)

Start loading the requested plans.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
35
36
def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
    """Start loading the requested plans."""

poll(job)

Return a finished job without blocking, or None if still pending.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
44
45
def poll(self, job: TransferJobState) -> TransferJobState | None:
    """Return a finished job without blocking, or None if still pending."""

publish(job)

Accept a completed store; a runtime may defer visibility until then.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
47
48
def publish(self, job: TransferJobState) -> None:
    """Accept a completed store; a runtime may defer visibility until then."""

register_buffers(kv_caches)

Register each unique KV allocation with the runtime.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
32
33
def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None:
    """Register each unique KV allocation with the runtime."""

store(plans)

Start storing the requested plans.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
38
39
def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
    """Start storing the requested plans."""

wait(job)

Block until the job is final and its buffers are safe to reuse.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py 
41
42
def wait(self, job: TransferJobState) -> TransferJobState:
    """Block until the job is final and its buffers are safe to reuse."""