Skip to content

vllm.distributed.ec_transfer.ec_connector.mooncake.metadata

Metadata exchanged by the Mooncake encoder-cache connector.

Classes:

ECMooncakeConnectorMetadata dataclass

Bases: ECConnectorMetadata

Worker operations emitted for one Scheduler step.

Source code in vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py 
45
46
47
48
49
50
51
52
@dataclass
class ECMooncakeConnectorMetadata(ECConnectorMetadata):
    """Worker operations emitted for one Scheduler step."""

    loads: list[ECMooncakeLoadSpec] = field(default_factory=list)
    pushes: list[ECMooncakePushSpec] = field(default_factory=list)
    freed: list[str] | None = None
    store_candidates: dict[str, TensorSpec] = field(default_factory=dict)

ECMooncakeLoadSpec dataclass

Describe a remote reservation or resident allocation to load.

Source code in vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py 
18
19
20
21
22
23
24
25
26
27
28
29
@dataclass
class ECMooncakeLoadSpec:
    """Describe a remote reservation or resident allocation to load."""

    mm_hash: str
    nbytes: int
    shape: tuple[int, ...]
    dtype: str
    transfer_id: str
    # The consumer pool still holds this item, so the load is a local handoff:
    # no transfer, no producer.
    local: bool = False

ECMooncakePushSpec dataclass

Describe a destination reservation prepared before a tensor is ready.

Source code in vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py 
32
33
34
35
36
37
38
39
40
41
42
@dataclass
class ECMooncakePushSpec:
    """Describe a destination reservation prepared before a tensor is ready."""

    mm_hash: str
    nbytes: int
    shape: tuple[int, ...]
    dtype: str
    consumer_zmq: str
    transfer_id: str
    request_id: str = ""

ECMooncakeWorkerMetadata dataclass

Bases: ECConnectorWorkerMetadata

Completion state reported from Workers to the Scheduler.

Source code in vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
@dataclass
class ECMooncakeWorkerMetadata(ECConnectorWorkerMetadata):
    """Completion state reported from Workers to the Scheduler."""

    loaded: set[str] = field(default_factory=set)
    failed_loads: set[str] = field(default_factory=set)
    # Items the receive pool dropped under pressure. The scheduler assumes an
    # evicted item stays resident until told otherwise.
    reclaimed: set[str] = field(default_factory=set)
    pending_saves: bool = False
    failed_saves: set[str] = field(default_factory=set)

    def aggregate(self, other: ECConnectorWorkerMetadata) -> ECMooncakeWorkerMetadata:
        assert isinstance(other, ECMooncakeWorkerMetadata)
        return ECMooncakeWorkerMetadata(
            # Every tensor-parallel rank gathers the embedding from its own
            # cache, so an item counts as loaded only where all of them have
            # it; one rank falling short must fail the load rather than leave
            # the scheduler believing it is ready.
            loaded=self.loaded & other.loaded,
            failed_loads=self.failed_loads | other.failed_loads,
            reclaimed=self.reclaimed | other.reclaimed,
            pending_saves=self.pending_saves or other.pending_saves,
            failed_saves=self.failed_saves | other.failed_saves,
        )