Skip to content

vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache

EmbeddingCache — named-entry block cache with FIFO eviction.

Manages a fixed pool of block IDs keyed by content identity (mm_hash). Entries transition through: not-ready → ready (evictable) → pinned. Eviction targets ready + unpinned entries in FIFO order.

All public methods are thread-safe.

Classes:

  • CacheEntry –

    A single cache entry. Read .block_ids and .ready freely;

  • EmbeddingCache –

    Fixed-size block cache with FIFO eviction.

CacheEntry

A single cache entry. Read .block_ids and .ready freely; mutations only through EmbeddingCache methods.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class CacheEntry:
    """A single cache entry. Read `.block_ids` and `.ready` freely;
    mutations only through EmbeddingCache methods."""

    __slots__ = ("block_ids", "_pin_count")

    def __init__(self, block_ids: tuple[int, ...]) -> None:
        self.block_ids = block_ids
        self._pin_count = -1  # not ready

    @property
    def ready(self) -> bool:
        return self._pin_count >= 0

    @property
    def evictable(self) -> bool:
        return self._pin_count == 0

    def mark_ready(self):
        self._pin_count = 0

    def pin(self):
        self._pin_count += 1

EmbeddingCache

Fixed-size block cache with FIFO eviction.

Entries are keyed by content identity (e.g. mm_hash). Blocks are allocated from a free-set; when space is needed, ready + unpinned entries are evicted oldest-first.

The caller is responsible for deciding when to call mark_ready (e.g. after enough engine steps have elapsed for the worker to have completed the write).

Methods:

  • alloc –

    Allocate n_blocks for key, evicting as needed.

  • discard –

    Remove a not-ready in-flight entry, returning its blocks to the pool.

  • get –

    Return the entry for key, or None if not present.

  • has_held_entries –

    True if any entry is not-ready or pinned (an in-flight transfer).

  • mark_ready –

    Mark an entry as ready (data is CPU-visible).

  • pin –

    Pin an entry (prevent eviction).

  • pin_if_ready –

    Atomically pin key if present and ready; return its entry.

  • unpin –

    Unpin an entry. Asserts currently pinned.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 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
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
class EmbeddingCache:
    """Fixed-size block cache with FIFO eviction.

    Entries are keyed by content identity (e.g. mm_hash). Blocks are
    allocated from a free-set; when space is needed, ready + unpinned
    entries are evicted oldest-first.

    The caller is responsible for deciding *when* to call ``mark_ready``
    (e.g. after enough engine steps have elapsed for the worker to have
    completed the write).
    """

    def __init__(self, num_blocks: int) -> None:
        self._num_blocks = num_blocks
        self._lock = threading.Lock()
        # Available block IDs, LIFO stack.
        self._free_blocks: list[int] = list(range(num_blocks))
        # Key → live entry (all states: not-ready, ready, pinned).
        self._entries: dict[str, CacheEntry] = {}
        # Ready + unpinned entries in insertion order (FIFO eviction).
        self._entries_free_list: OrderedDict[str, None] = OrderedDict()
        # Total blocks reclaimable via eviction (fail-fast check).
        self._evictable_block_count: int = 0

    def get(self, key: str) -> CacheEntry | None:
        """Return the entry for *key*, or None if not present."""
        with self._lock:
            return self._entries.get(key)

    def has_held_entries(self) -> bool:
        """True if any entry is not-ready or pinned (an in-flight transfer).

        Held entries are exactly those absent from the eviction free list, so
        they represent saves awaiting mark_ready or loads awaiting unpin.
        """
        with self._lock:
            return len(self._entries) != len(self._entries_free_list)

    def alloc(self, key: str, n_blocks: int) -> CacheEntry | None:
        """Allocate *n_blocks* for *key*, evicting as needed.

        The entry starts not-ready. Returns None if there is not enough
        space even after evicting all evictable entries.
        """
        with self._lock:
            assert key not in self._entries, (
                f"EmbeddingCache: duplicate alloc for {key!r}"
            )
            assert n_blocks <= self._num_blocks, (
                f"EmbeddingCache: {n_blocks} blocks requested but capacity "
                f"is {self._num_blocks}"
            )
            if len(self._free_blocks) + self._evictable_block_count < n_blocks:
                return None
            self._evict_until(n_blocks)
            # Canonical ascending order: keeps a freed entry's blocks contiguous
            # in the free list when it is evicted, and lets consumers detect
            # which of an entry's blocks are adjacent.
            block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks)))
            entry = CacheEntry(block_ids)
            self._entries[key] = entry
            return entry

    def mark_ready(self, key: str) -> None:
        """Mark an entry as ready (data is CPU-visible)."""
        with self._lock:
            entry = self._entries[key]
            assert entry._pin_count == -1, (
                f"EmbeddingCache: mark_ready on already-ready entry {key!r}"
            )
            entry.mark_ready()
            self._entries_free_list[key] = None
            self._evictable_block_count += len(entry.block_ids)

    def pin(self, key: str) -> None:
        """Pin an entry (prevent eviction)."""
        with self._lock:
            entry = self._entries[key]
            assert entry.ready, f"EmbeddingCache: pin of not-ready entry {key!r}"
            if entry.evictable:
                del self._entries_free_list[key]
                self._evictable_block_count -= len(entry.block_ids)
            entry.pin()

    def pin_if_ready(self, key: str) -> CacheEntry | None:
        """Atomically pin *key* if present and ready; return its entry.

        Returns None when the key is absent. An entry returned with `ready`
        False was not pinned: its save is still in flight, which the producer
        reports differently from a miss because such a read can be retried.
        """
        with self._lock:
            entry = self._entries.get(key)
            if entry is None or not entry.ready:
                return entry
            if entry.evictable:
                del self._entries_free_list[key]
                self._evictable_block_count -= len(entry.block_ids)
            entry.pin()
            return entry

    def unpin(self, key: str) -> None:
        """Unpin an entry. Asserts currently pinned."""
        with self._lock:
            entry = self._entries[key]
            assert entry._pin_count > 0, (
                f"EmbeddingCache: unpin of unpinned entry {key!r}"
            )
            entry._pin_count -= 1
            if entry._pin_count == 0:
                self._entries_free_list[key] = None
                self._evictable_block_count += len(entry.block_ids)

    def discard(self, key: str) -> None:
        """Remove a not-ready in-flight entry, returning its blocks to the pool.

        Used when an in-flight fill (e.g. a NIXL READ) fails before the entry
        is marked ready. Asserts the entry is present and not ready — ready or
        pinned entries are reclaimed through eviction/unpin, not discard.
        """
        with self._lock:
            entry = self._entries[key]
            assert entry._pin_count == -1, (
                f"EmbeddingCache: discard of ready/pinned entry {key!r}"
            )
            del self._entries[key]
            self._free_blocks.extend(entry.block_ids)

    def _evict_until(self, n_blocks: int) -> None:
        """Evict ready+unpinned entries FIFO until enough space. Lock held."""
        while len(self._free_blocks) < n_blocks:
            key, _ = self._entries_free_list.popitem(last=False)
            entry = self._entries.pop(key)
            self._evictable_block_count -= len(entry.block_ids)
            self._free_blocks.extend(entry.block_ids)

_evict_until(n_blocks)

Evict ready+unpinned entries FIFO until enough space. Lock held.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
169
170
171
172
173
174
175
def _evict_until(self, n_blocks: int) -> None:
    """Evict ready+unpinned entries FIFO until enough space. Lock held."""
    while len(self._free_blocks) < n_blocks:
        key, _ = self._entries_free_list.popitem(last=False)
        entry = self._entries.pop(key)
        self._evictable_block_count -= len(entry.block_ids)
        self._free_blocks.extend(entry.block_ids)

alloc(key, n_blocks)

Allocate n_blocks for key, evicting as needed.

The entry starts not-ready. Returns None if there is not enough space even after evicting all evictable entries.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def alloc(self, key: str, n_blocks: int) -> CacheEntry | None:
    """Allocate *n_blocks* for *key*, evicting as needed.

    The entry starts not-ready. Returns None if there is not enough
    space even after evicting all evictable entries.
    """
    with self._lock:
        assert key not in self._entries, (
            f"EmbeddingCache: duplicate alloc for {key!r}"
        )
        assert n_blocks <= self._num_blocks, (
            f"EmbeddingCache: {n_blocks} blocks requested but capacity "
            f"is {self._num_blocks}"
        )
        if len(self._free_blocks) + self._evictable_block_count < n_blocks:
            return None
        self._evict_until(n_blocks)
        # Canonical ascending order: keeps a freed entry's blocks contiguous
        # in the free list when it is evicted, and lets consumers detect
        # which of an entry's blocks are adjacent.
        block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks)))
        entry = CacheEntry(block_ids)
        self._entries[key] = entry
        return entry

discard(key)

Remove a not-ready in-flight entry, returning its blocks to the pool.

Used when an in-flight fill (e.g. a NIXL READ) fails before the entry is marked ready. Asserts the entry is present and not ready — ready or pinned entries are reclaimed through eviction/unpin, not discard.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def discard(self, key: str) -> None:
    """Remove a not-ready in-flight entry, returning its blocks to the pool.

    Used when an in-flight fill (e.g. a NIXL READ) fails before the entry
    is marked ready. Asserts the entry is present and not ready — ready or
    pinned entries are reclaimed through eviction/unpin, not discard.
    """
    with self._lock:
        entry = self._entries[key]
        assert entry._pin_count == -1, (
            f"EmbeddingCache: discard of ready/pinned entry {key!r}"
        )
        del self._entries[key]
        self._free_blocks.extend(entry.block_ids)

get(key)

Return the entry for key, or None if not present.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
65
66
67
68
def get(self, key: str) -> CacheEntry | None:
    """Return the entry for *key*, or None if not present."""
    with self._lock:
        return self._entries.get(key)

has_held_entries()

True if any entry is not-ready or pinned (an in-flight transfer).

Held entries are exactly those absent from the eviction free list, so they represent saves awaiting mark_ready or loads awaiting unpin.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
70
71
72
73
74
75
76
77
def has_held_entries(self) -> bool:
    """True if any entry is not-ready or pinned (an in-flight transfer).

    Held entries are exactly those absent from the eviction free list, so
    they represent saves awaiting mark_ready or loads awaiting unpin.
    """
    with self._lock:
        return len(self._entries) != len(self._entries_free_list)

mark_ready(key)

Mark an entry as ready (data is CPU-visible).

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
104
105
106
107
108
109
110
111
112
113
def mark_ready(self, key: str) -> None:
    """Mark an entry as ready (data is CPU-visible)."""
    with self._lock:
        entry = self._entries[key]
        assert entry._pin_count == -1, (
            f"EmbeddingCache: mark_ready on already-ready entry {key!r}"
        )
        entry.mark_ready()
        self._entries_free_list[key] = None
        self._evictable_block_count += len(entry.block_ids)

pin(key)

Pin an entry (prevent eviction).

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
115
116
117
118
119
120
121
122
123
def pin(self, key: str) -> None:
    """Pin an entry (prevent eviction)."""
    with self._lock:
        entry = self._entries[key]
        assert entry.ready, f"EmbeddingCache: pin of not-ready entry {key!r}"
        if entry.evictable:
            del self._entries_free_list[key]
            self._evictable_block_count -= len(entry.block_ids)
        entry.pin()

pin_if_ready(key)

Atomically pin key if present and ready; return its entry.

Returns None when the key is absent. An entry returned with ready False was not pinned: its save is still in flight, which the producer reports differently from a miss because such a read can be retried.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def pin_if_ready(self, key: str) -> CacheEntry | None:
    """Atomically pin *key* if present and ready; return its entry.

    Returns None when the key is absent. An entry returned with `ready`
    False was not pinned: its save is still in flight, which the producer
    reports differently from a miss because such a read can be retried.
    """
    with self._lock:
        entry = self._entries.get(key)
        if entry is None or not entry.ready:
            return entry
        if entry.evictable:
            del self._entries_free_list[key]
            self._evictable_block_count -= len(entry.block_ids)
        entry.pin()
        return entry

unpin(key)

Unpin an entry. Asserts currently pinned.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py 
142
143
144
145
146
147
148
149
150
151
152
def unpin(self, key: str) -> None:
    """Unpin an entry. Asserts currently pinned."""
    with self._lock:
        entry = self._entries[key]
        assert entry._pin_count > 0, (
            f"EmbeddingCache: unpin of unpinned entry {key!r}"
        )
        entry._pin_count -= 1
        if entry._pin_count == 0:
            self._entries_free_list[key] = None
            self._evictable_block_count += len(entry.block_ids)