@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 43c8cbf79b

[EC Connector] CPU Offloading EC Connector (#47423)

Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

omerpaz95 committed Jul 13, 2026 at 20:09 UTC 43c8cbf79b75a583ef870b13e82dd33e68cf32b1
18 files changed +3008
tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py new
+306
@@ -0,0 +1,306 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for EmbeddingCache.
4 +
5 +Pure integer bookkeeping — no mmap, torch, or CUDA — so these run anywhere.
6 +"""
7 +
8 +import threading
9 +
10 +import pytest
11 +
12 +from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
13 + EmbeddingCache,
14 +)
15 +
16 +
17 +def _cache(num_blocks: int = 8) -> EmbeddingCache:
18 + return EmbeddingCache(num_blocks)
19 +
20 +
21 +# ── alloc ────────────────────────────────────────────────────────────────────
22 +
23 +
24 +def test_alloc_returns_unique_block_ids():
25 + cache = _cache()
26 + entry = cache.alloc("a", 3)
27 + assert entry is not None
28 + assert len(entry.block_ids) == 3
29 + assert len(set(entry.block_ids)) == 3
30 + assert all(0 <= i < 8 for i in entry.block_ids)
31 +
32 +
33 +def test_alloc_returns_tuple():
34 + cache = _cache()
35 + entry = cache.alloc("a", 2)
36 + assert isinstance(entry.block_ids, tuple)
37 +
38 +
39 +def test_alloc_returns_none_when_full():
40 + cache = _cache(num_blocks=4)
41 + cache.alloc("a", 4)
42 + assert cache.alloc("b", 1) is None
43 +
44 +
45 +def test_alloc_asserts_duplicate_key():
46 + cache = _cache()
47 + cache.alloc("a", 2)
48 + with pytest.raises(AssertionError, match="duplicate alloc"):
49 + cache.alloc("a", 2)
50 +
51 +
52 +def test_alloc_asserts_n_blocks_exceeds_capacity():
53 + cache = _cache(num_blocks=4)
54 + with pytest.raises(AssertionError, match="capacity"):
55 + cache.alloc("a", 5)
56 +
57 +
58 +def test_alloc_zero_blocks():
59 + cache = _cache()
60 + entry = cache.alloc("a", 0)
61 + assert entry is not None
62 + assert entry.block_ids == ()
63 +
64 +
65 +def test_alloc_starts_not_ready():
66 + cache = _cache()
67 + entry = cache.alloc("a", 2)
68 + assert not entry.ready
69 + assert not entry.evictable
70 +
71 +
72 +# ── get ──────────────────────────────────────────────────────────────────────
73 +
74 +
75 +def test_get_missing_returns_none():
76 + cache = _cache()
77 + assert cache.get("a") is None
78 +
79 +
80 +def test_get_returns_entry():
81 + cache = _cache()
82 + alloc_entry = cache.alloc("a", 2)
83 + got = cache.get("a")
84 + assert got is alloc_entry
85 +
86 +
87 +def test_get_returns_not_ready_entry():
88 + cache = _cache()
89 + cache.alloc("a", 2)
90 + entry = cache.get("a")
91 + assert entry is not None
92 + assert not entry.ready
93 +
94 +
95 +# ── mark_ready ───────────────────────────────────────────────────────────────
96 +
97 +
98 +def test_mark_ready_transitions():
99 + cache = _cache()
100 + cache.alloc("a", 2)
101 + entry = cache.get("a")
102 + assert not entry.ready
103 + cache.mark_ready("a")
104 + assert entry.ready
105 + assert entry.evictable
106 +
107 +
108 +def test_mark_ready_asserts_already_ready():
109 + cache = _cache()
110 + cache.alloc("a", 2)
111 + cache.mark_ready("a")
112 + with pytest.raises(AssertionError, match="already-ready"):
113 + cache.mark_ready("a")
114 +
115 +
116 +def test_mark_ready_asserts_missing_key():
117 + cache = _cache()
118 + with pytest.raises(KeyError):
119 + cache.mark_ready("nonexistent")
120 +
121 +
122 +# ── pin / unpin ──────────────────────────────────────────────────────────────
123 +
124 +
125 +def test_pin_asserts_not_ready():
126 + cache = _cache()
127 + cache.alloc("a", 2)
128 + with pytest.raises(AssertionError, match="not-ready"):
129 + cache.pin("a")
130 +
131 +
132 +def test_pin_unpin_basic():
133 + cache = _cache()
134 + cache.alloc("a", 2)
135 + cache.mark_ready("a")
136 + cache.pin("a")
137 + entry = cache.get("a")
138 + assert not entry.evictable
139 + cache.unpin("a")
140 + assert entry.evictable
141 +
142 +
143 +def test_nested_pin():
144 + cache = _cache()
145 + cache.alloc("a", 2)
146 + cache.mark_ready("a")
147 + cache.pin("a")
148 + cache.pin("a")
149 + cache.unpin("a")
150 + entry = cache.get("a")
151 + assert not entry.evictable # still pinned (count=1)
152 + cache.unpin("a")
153 + assert entry.evictable
154 +
155 +
156 +def test_unpin_asserts_on_unpinned():
157 + cache = _cache()
158 + cache.alloc("a", 2)
159 + cache.mark_ready("a")
160 + with pytest.raises(AssertionError, match="unpinned"):
161 + cache.unpin("a")
162 +
163 +
164 +# ── eviction ─────────────────────────────────────────────────────────────────
165 +
166 +
167 +def test_eviction_frees_ready_unpinned_fifo():
168 + cache = _cache(num_blocks=4)
169 + cache.alloc("a", 2)
170 + cache.mark_ready("a")
171 + cache.alloc("b", 2)
172 + cache.mark_ready("b")
173 + # Pool full. Alloc should evict "a" (oldest).
174 + entry_c = cache.alloc("c", 2)
175 + assert entry_c is not None
176 + assert cache.get("a") is None # evicted
177 + assert cache.get("b") is not None # still there
178 +
179 +
180 +def test_eviction_skips_pinned():
181 + cache = _cache(num_blocks=4)
182 + cache.alloc("a", 2)
183 + cache.mark_ready("a")
184 + cache.pin("a")
185 + cache.alloc("b", 2)
186 + cache.mark_ready("b")
187 + # "a" is pinned, so "b" should be evicted instead.
188 + entry_c = cache.alloc("c", 2)
189 + assert entry_c is not None
190 + assert cache.get("a") is not None # pinned, survived
191 + assert cache.get("b") is None # evicted
192 + cache.unpin("a")
193 +
194 +
195 +def test_eviction_skips_not_ready():
196 + cache = _cache(num_blocks=4)
197 + cache.alloc("a", 2) # not ready
198 + cache.alloc("b", 2)
199 + cache.mark_ready("b")
200 + # "a" is not ready (non-evictable), "b" is evictable.
201 + entry_c = cache.alloc("c", 2)
202 + assert entry_c is not None
203 + assert cache.get("a") is not None # not ready, survived
204 + assert cache.get("b") is None # evicted
205 + cache.mark_ready("a") # clean up
206 +
207 +
208 +def test_eviction_returns_none_if_all_non_evictable():
209 + cache = _cache(num_blocks=4)
210 + cache.alloc("a", 2)
211 + cache.mark_ready("a")
212 + cache.pin("a")
213 + cache.alloc("b", 2) # not ready
214 + # Both entries are non-evictable (a=pinned, b=not-ready).
215 + assert cache.alloc("c", 1) is None
216 + # Fail-fast: nothing was evicted.
217 + assert cache.get("a") is not None
218 + assert cache.get("b") is not None
219 + cache.unpin("a")
220 +
221 +
222 +def test_fail_fast_no_wasted_eviction():
223 + cache = _cache(num_blocks=6)
224 + cache.alloc("a", 2)
225 + cache.mark_ready("a") # evictable: 2 blocks
226 + cache.alloc("b", 2) # not ready
227 + cache.alloc("c", 2) # not ready
228 + # free=0, evictable=2, need 3 -> cannot succeed
229 + assert cache.alloc("d", 3) is None
230 + # "a" was NOT evicted despite being evictable
231 + assert cache.get("a") is not None
232 +
233 +
234 +def test_eviction_multiple_entries():
235 + cache = _cache(num_blocks=6)
236 + cache.alloc("a", 2)
237 + cache.mark_ready("a")
238 + cache.alloc("b", 2)
239 + cache.mark_ready("b")
240 + cache.alloc("c", 2)
241 + cache.mark_ready("c")
242 + # Need 4 blocks: must evict both "a" and "b".
243 + entry_d = cache.alloc("d", 4)
244 + assert entry_d is not None
245 + assert cache.get("a") is None
246 + assert cache.get("b") is None
247 + assert cache.get("c") is not None
248 +
249 +
250 +# ── thread safety ────────────────────────────────────────────────────────────
251 +
252 +
253 +def test_concurrent_pin_unpin():
254 + cache = _cache(num_blocks=8)
255 + cache.alloc("a", 4)
256 + cache.mark_ready("a")
257 + errors: list[Exception] = []
258 + barrier = threading.Barrier(8)
259 +
260 + def worker():
261 + barrier.wait()
262 + try:
263 + for _ in range(500):
264 + cache.pin("a")
265 + cache.unpin("a")
266 + except Exception as exc:
267 + errors.append(exc)
268 +
269 + threads = [threading.Thread(target=worker) for _ in range(8)]
270 + for t in threads:
271 + t.start()
272 + for t in threads:
273 + t.join()
274 +
275 + assert not errors, errors
276 + entry = cache.get("a")
277 + assert entry.evictable # all pins balanced
278 +
279 +
280 +def test_concurrent_alloc_evict():
281 + cache = _cache(num_blocks=4)
282 + errors: list[Exception] = []
283 + barrier = threading.Barrier(4)
284 + counter = {"i": 0}
285 + counter_lock = threading.Lock()
286 +
287 + def worker():
288 + barrier.wait()
289 + try:
290 + for _ in range(200):
291 + with counter_lock:
292 + key = f"k{counter['i']}"
293 + counter["i"] += 1
294 + entry = cache.alloc(key, 1)
295 + if entry is not None:
296 + cache.mark_ready(key)
297 + except Exception as exc:
298 + errors.append(exc)
299 +
300 + threads = [threading.Thread(target=worker) for _ in range(4)]
301 + for t in threads:
302 + t.start()
303 + for t in threads:
304 + t.join()
305 +
306 + assert not errors, errors
tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py new
+480
@@ -0,0 +1,480 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import uuid
4 +
5 +import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
6 +from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
7 + ECSharedRegion,
8 +)
9 +from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
10 +
11 +_N = 16
12 +_BS = 64
13 +
14 +
15 +class _Pos:
16 + def __init__(self, offset, length):
17 + self.offset = offset
18 + self.length = length
19 +
20 +
21 +class _Feature:
22 + def __init__(self, mm_hash, length=1, identifier=None):
23 + self.mm_hash = mm_hash
24 + self.identifier = identifier if identifier is not None else mm_hash
25 + self.mm_position = _Pos(0, length)
26 +
27 +
28 +class _Request:
29 + _counter = 0
30 +
31 + def __init__(self, features, request_id=None):
32 + self.mm_features = features
33 + if request_id is None:
34 + _Request._counter += 1
35 + self.request_id = f"req_{_Request._counter}"
36 + else:
37 + self.request_id = request_id
38 +
39 +
40 +def _make_scheduler(
41 + monkeypatch,
42 + num_blocks=_N,
43 + max_concurrent_batches=1,
44 + *,
45 + is_producer=True,
46 + is_consumer=True,
47 +) -> ECCPUScheduler:
48 + region = ECSharedRegion(
49 + engine_id=str(uuid.uuid4()),
50 + num_blocks=num_blocks,
51 + block_size_bytes=_BS,
52 + )
53 + monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)
54 +
55 + _is_prod = is_producer
56 + _is_cons = is_consumer
57 + _mcb = max_concurrent_batches
58 +
59 + class _EC:
60 + is_ec_producer = _is_prod
61 + is_ec_consumer = _is_cons
62 +
63 + class _Cfg:
64 + ec_transfer_config = _EC()
65 +
66 + @property
67 + def max_concurrent_batches(self):
68 + return _mcb
69 +
70 + return ECCPUScheduler(_Cfg())
71 +
72 +
73 +class _SchedulerOutput:
74 + def __init__(self, finished_req_ids=None):
75 + self.finished_req_ids = finished_req_ids or set()
76 +
77 +
78 +def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int):
79 + """Pre-populate a ready cache entry backed by real blocks."""
80 + s._cache.alloc(mm_hash, n_blocks)
81 + s._cache.mark_ready(mm_hash)
82 +
83 +
84 +def test_offload_reuse_cycle(monkeypatch):
85 + s = _make_scheduler(monkeypatch)
86 + req = _Request([_Feature("h1", length=1)])
87 +
88 + # Step A: first sight — allocate and save.
89 + assert s.has_cache_item("h1") is False
90 + assert s.ensure_cache_available(req, 0) is True
91 + s.update_state_after_alloc(req, 0)
92 + meta_a = s.build_connector_meta(scheduler_output=None)
93 + assert "h1" in meta_a.saves
94 + assert meta_a.loads == {}
95 + # Not ready yet (max_concurrent_batches=1 means ready next step).
96 + assert s.has_cache_item("h1") is False
97 +
98 + # Step B: _step_readiness marks h1 ready (no new allocs).
99 + s.build_connector_meta(scheduler_output=None)
100 + assert s.has_cache_item("h1") is True
101 +
102 + # Step C: reload now works — entry is ready at update_state_after_alloc time.
103 + s.update_state_after_alloc(req, 0)
104 + meta_c = s.build_connector_meta(scheduler_output=None)
105 + assert "h1" in meta_c.loads
106 + assert meta_c.loads["h1"] == meta_a.saves["h1"]
107 +
108 + s.shutdown()
109 +
110 +
111 +def test_has_cache_item_false_when_not_consumer(monkeypatch):
112 + s = _make_scheduler(monkeypatch, is_consumer=False)
113 + assert s.has_cache_item("anything") is False
114 + s.shutdown()
115 +
116 +
117 +def test_connector_keys_on_identifier_not_mm_hash(monkeypatch):
118 + """The connector must key the encoder cache on feature.identifier (what
119 + has_cache_item is called with), NOT feature.mm_hash."""
120 + s = _make_scheduler(monkeypatch)
121 + req = _Request([_Feature("PROC_KEY", length=1, identifier="ENC_KEY")])
122 + s.update_state_after_alloc(req, 0)
123 + meta = s.build_connector_meta(scheduler_output=None)
124 + assert "ENC_KEY" in meta.saves
125 + assert "PROC_KEY" not in meta.saves
126 + # Not ready yet (1 step delay).
127 + assert s.has_cache_item("ENC_KEY") is False
128 + # After another step it becomes ready — keyed on identifier.
129 + s.build_connector_meta(scheduler_output=None)
130 + assert s.has_cache_item("ENC_KEY") is True
131 + assert s.has_cache_item("PROC_KEY") is False
132 + s.shutdown()
133 +
134 +
135 +def test_cpu_region_fifo_eviction(monkeypatch):
136 + # Region holds exactly 2 one-block encodings.
137 + s = _make_scheduler(monkeypatch, num_blocks=2)
138 +
139 + # Fill with A and B, marking each ready.
140 + for h in ("A", "B"):
141 + s.update_state_after_alloc(_Request([_Feature(h, length=1)]), 0)
142 + s.build_connector_meta(scheduler_output=None)
143 + # After the step above, A's step was called in B's build. A is ready.
144 + # B was submitted in the last build, so one more step makes it ready.
145 + s.build_connector_meta(scheduler_output=None) # marks B ready
146 + assert s.has_cache_item("A") is True
147 + assert s.has_cache_item("B") is True
148 +
149 + # Allocate C: must evict A (oldest) to make room.
150 + s.update_state_after_alloc(_Request([_Feature("C", length=1)]), 0)
151 + s.build_connector_meta(scheduler_output=None)
152 + assert s.has_cache_item("A") is False # evicted
153 + assert s.has_cache_item("B") is True
154 +
155 +
156 +def test_multiple_mm_items_per_request(monkeypatch):
157 + s = _make_scheduler(monkeypatch)
158 + req = _Request([_Feature("h1", length=1), _Feature("h2", length=1)])
159 + s.update_state_after_alloc(req, 0)
160 + s.update_state_after_alloc(req, 1)
161 + meta = s.build_connector_meta(scheduler_output=None)
162 + assert set(meta.saves) == {"h1", "h2"}
163 +
164 + # Step delay: mark both ready.
165 + s.build_connector_meta(scheduler_output=None)
166 + assert s.has_cache_item("h1") is True
167 + assert s.has_cache_item("h2") is True
168 +
169 + # Now reload both.
170 + s.update_state_after_alloc(req, 0)
171 + s.update_state_after_alloc(req, 1)
172 + meta2 = s.build_connector_meta(scheduler_output=None)
173 + assert set(meta2.loads) == {"h1", "h2"}
174 + s.shutdown()
175 +
176 +
177 +def test_load_returns_correct_block_ids(monkeypatch):
178 + """meta.loads must contain the same block IDs that meta.saves
179 + allocated — verified through the public API only."""
180 + s = _make_scheduler(monkeypatch)
181 + req = _Request([_Feature("a", length=2)])
182 +
183 + s.update_state_after_alloc(req, 0)
184 + meta_save = s.build_connector_meta(scheduler_output=None)
185 + assert "a" in meta_save.saves
186 +
187 + # Readiness delay.
188 + s.build_connector_meta(scheduler_output=None)
189 +
190 + # Reload.
191 + s.update_state_after_alloc(req, 0)
192 + meta_load = s.build_connector_meta(scheduler_output=None)
193 + assert meta_load.loads["a"] == meta_save.saves["a"]
194 + s.shutdown()
195 +
196 +
197 +def test_loads_only_serves_hashes_touched_this_step(monkeypatch):
198 + s = _make_scheduler(monkeypatch)
199 + _seed_cached(s, "a", n_blocks=2)
200 + _seed_cached(s, "b", n_blocks=2)
201 +
202 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)
203 + meta = s.build_connector_meta(scheduler_output=None)
204 + assert "a" in meta.loads
205 + assert "b" not in meta.loads
206 + s.shutdown()
207 +
208 +
209 +def test_repeated_reload_same_step_loads_once(monkeypatch):
210 + """The same mm_hash requested twice in one step must appear in
211 + meta.loads exactly once."""
212 + s = _make_scheduler(monkeypatch)
213 + _seed_cached(s, "a", n_blocks=2)
214 +
215 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)
216 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)
217 + meta = s.build_connector_meta(scheduler_output=None)
218 + assert list(meta.loads.keys()).count("a") == 1
219 + s.shutdown()
220 +
221 +
222 +def test_load_not_emitted_for_uncached_entry(monkeypatch):
223 + """A feature that was never saved must not appear in meta.loads."""
224 + s = _make_scheduler(monkeypatch)
225 + s.update_state_after_alloc(_Request([_Feature("missing")]), 0)
226 + meta = s.build_connector_meta(scheduler_output=None)
227 + assert meta.loads == {}
228 + s.shutdown()
229 +
230 +
231 +def test_delayed_unpin_protects_blocks_during_worker_read(monkeypatch):
232 + """Blocks being loaded must not be evictable until the delayed unpin
233 + fires. Observable: a save that needs eviction cannot reclaim a
234 + still-pinned loaded entry."""
235 + # 2 blocks, both occupied by "a" (ready).
236 + s = _make_scheduler(monkeypatch, num_blocks=2)
237 + _seed_cached(s, "a", n_blocks=2)
238 +
239 + # Load "a" — pins it.
240 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)
241 + meta = s.build_connector_meta(scheduler_output=None)
242 + assert "a" in meta.loads
243 +
244 + # Try to save "b" (needs 2 blocks, but "a" is pinned) — must fail.
245 + s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
246 + meta2 = s.build_connector_meta(scheduler_output=None)
247 + assert "b" not in meta2.saves # can't evict pinned "a"
248 +
249 + # After the delay, "a" unpins. Now "b" can evict it.
250 + s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
251 + meta3 = s.build_connector_meta(scheduler_output=None)
252 + assert "b" in meta3.saves
253 + s.shutdown()
254 +
255 +
256 +def test_delayed_unpin_depth_2(monkeypatch):
257 + """With max_concurrent_batches=2, loaded blocks stay pinned for 2 steps
258 + after build. Observable via eviction failure."""
259 + s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=2)
260 + _seed_cached(s, "a", n_blocks=2)
261 +
262 + # Load "a" → pinned.
263 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)
264 + s.build_connector_meta(scheduler_output=None)
265 +
266 + # Step 1 after load: still pinned — can't evict for "b".
267 + s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
268 + meta1 = s.build_connector_meta(scheduler_output=None)
269 + assert "b" not in meta1.saves
270 +
271 + # Step 2 after load: still pinned.
272 + s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
273 + meta2 = s.build_connector_meta(scheduler_output=None)
274 + assert "b" not in meta2.saves
275 +
276 + # Step 3: unpin fires, now eviction works.
277 + s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
278 + meta3 = s.build_connector_meta(scheduler_output=None)
279 + assert "b" in meta3.saves
280 + s.shutdown()
281 +
282 +
283 +def test_shutdown_unpins_pending_loads(monkeypatch):
284 + """Shutdown must unpin entries still in _pending_loads (not yet drained
285 + to the unpin deque). Observable: after shutdown, the entry is evictable
286 + (alloc can reclaim its blocks)."""
287 + s = _make_scheduler(monkeypatch, num_blocks=2)
288 + _seed_cached(s, "a", n_blocks=2)
289 +
290 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)
291 + # Don't call build_connector_meta — "a" stays in _pending_loads.
292 + s.shutdown()
293 +
294 + # After shutdown, scheduler reports no items (consumer role is off).
295 + assert s.has_cache_item("a") is False
296 +
297 +
298 +def test_shutdown_unpins_deferred_unpin_queue(monkeypatch):
299 + """Shutdown must unpin entries queued in the deferred unpin deque
300 + (already drained from _pending_loads but not yet expired)."""
301 + s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=2)
302 + _seed_cached(s, "a", n_blocks=2)
303 +
304 + # Load "a" and build — moves to unpin deque.
305 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)
306 + s.build_connector_meta(scheduler_output=None)
307 +
308 + # "a" is in the deferred unpin deque (depth=2, not yet expired).
309 + # Shutdown must unpin it.
310 + s.shutdown()
311 + assert s.has_cache_item("a") is False
312 +
313 +
314 +def test_eviction_skips_entry_pinned_by_pending_load(monkeypatch):
315 + """A loaded entry still within the unpin delay window must survive
316 + eviction attempts from new saves."""
317 + # 2 blocks: A occupies both, ready.
318 + s = _make_scheduler(monkeypatch, num_blocks=2)
319 + _seed_cached(s, "A", n_blocks=2)
320 +
321 + # Load A → pinned.
322 + s.update_state_after_alloc(_Request([_Feature("A")]), 0)
323 + meta_load = s.build_connector_meta(scheduler_output=None)
324 + assert "A" in meta_load.loads
325 +
326 + # Try to save C (needs 2 blocks) — can't evict pinned A.
327 + s.update_state_after_alloc(_Request([_Feature("C", length=2)]), 0)
328 + meta_save = s.build_connector_meta(scheduler_output=None)
329 + assert "C" not in meta_save.saves
330 + assert s.has_cache_item("A") is True # survived
331 + s.shutdown()
332 +
333 +
334 +def test_region_full_skips_save_and_never_blocks(monkeypatch):
335 + """When the region is fully occupied by pinned entries, new saves are
336 + silently skipped and ensure_cache_available never blocks."""
337 + s = _make_scheduler(monkeypatch, num_blocks=1)
338 + _seed_cached(s, "pinned", n_blocks=1)
339 + s._cache.pin("pinned")
340 +
341 + req = _Request([_Feature("new", length=1)])
342 + assert s.ensure_cache_available(req, 0) is True
343 +
344 + s.update_state_after_alloc(req, 0)
345 + meta = s.build_connector_meta(scheduler_output=None)
346 + assert "new" not in meta.saves
347 + assert s.has_cache_item("new") is False
348 +
349 + s._cache.unpin("pinned")
350 + s.shutdown()
351 +
352 +
353 +def test_step_readiness_depth_2(monkeypatch):
354 + """With max_concurrent_batches=2, entries need 2 steps to become ready."""
355 + s = _make_scheduler(monkeypatch, max_concurrent_batches=2)
356 + req = _Request([_Feature("h1", length=1)])
357 +
358 + s.update_state_after_alloc(req, 0)
359 + meta_0 = s.build_connector_meta(scheduler_output=None)
360 + assert "h1" in meta_0.saves
361 + assert s.has_cache_item("h1") is False
362 +
363 + # One step later: still not ready (depth=2).
364 + s.build_connector_meta(scheduler_output=None)
365 + assert s.has_cache_item("h1") is False
366 +
367 + # Two steps later: now ready.
368 + s.build_connector_meta(scheduler_output=None)
369 + assert s.has_cache_item("h1") is True
370 + s.shutdown()
371 +
372 +
373 +def test_producer_only_never_emits_loads(monkeypatch):
374 + """A producer-only scheduler must never populate meta.loads, even when
375 + entries are ready."""
376 + s = _make_scheduler(monkeypatch, is_consumer=False)
377 + req = _Request([_Feature("h1", length=1)])
378 +
379 + s.update_state_after_alloc(req, 0)
380 + meta = s.build_connector_meta(scheduler_output=None)
381 + assert "h1" in meta.saves
382 + assert meta.loads == {}
383 +
384 + # Mark ready.
385 + s.build_connector_meta(scheduler_output=None)
386 +
387 + # Even if we re-encounter the same feature, no load.
388 + s.update_state_after_alloc(req, 0)
389 + meta2 = s.build_connector_meta(scheduler_output=None)
390 + assert meta2.loads == {}
391 + s.shutdown()
392 +
393 +
394 +def test_consumer_only_never_emits_saves(monkeypatch):
395 + """A consumer-only scheduler must never populate meta.saves."""
396 + s = _make_scheduler(monkeypatch, is_producer=False)
397 + _seed_cached(s, "a", n_blocks=1)
398 +
399 + req = _Request([_Feature("a", length=1)])
400 + s.update_state_after_alloc(req, 0)
401 + meta = s.build_connector_meta(scheduler_output=None)
402 + assert meta.saves == {}
403 + assert "a" in meta.loads
404 + s.shutdown()
405 +
406 +
407 +def test_consumer_only_has_cache_item(monkeypatch):
408 + """Consumer-only: has_cache_item works normally for ready entries."""
409 + s = _make_scheduler(monkeypatch, is_producer=False)
410 + assert s.has_cache_item("a") is False
411 + _seed_cached(s, "a", n_blocks=1)
412 + assert s.has_cache_item("a") is True
413 + s.shutdown()
414 +
415 +
416 +def test_save_not_emitted_for_already_cached_entry(monkeypatch):
417 + """An entry already in the cache (from a prior save) must not trigger
418 + a second allocation or appear in meta.saves again."""
419 + s = _make_scheduler(monkeypatch)
420 + req = _Request([_Feature("h1", length=1)])
421 +
422 + # First encounter → save.
423 + s.update_state_after_alloc(req, 0)
424 + meta1 = s.build_connector_meta(scheduler_output=None)
425 + assert "h1" in meta1.saves
426 +
427 + # Second encounter (same step or later) → no save.
428 + s.update_state_after_alloc(req, 0)
429 + meta2 = s.build_connector_meta(scheduler_output=None)
430 + assert "h1" not in meta2.saves
431 + s.shutdown()
432 +
433 +
434 +# ── first-finish fast-path (idle engine) ────────────────────────────────────
435 +
436 +
437 +def test_first_finish_marks_ready_immediately(monkeypatch):
438 + """When the originating request finishes, the save entry becomes ready
439 + without waiting for the step-count delay."""
440 + s = _make_scheduler(monkeypatch, max_concurrent_batches=5)
441 + req = _Request([_Feature("h1", length=1)], request_id="req_x")
442 +
443 + s.update_state_after_alloc(req, 0)
444 + s.build_connector_meta(scheduler_output=None)
445 + assert s.has_cache_item("h1") is False
446 +
447 + # Request finishes — readiness should fire immediately.
448 + s.build_connector_meta(_SchedulerOutput(finished_req_ids={"req_x"}))
449 + assert s.has_cache_item("h1") is True
450 + s.shutdown()
451 +
452 +
453 +def test_first_finish_unpins_immediately(monkeypatch):
454 + """When the originating request finishes, the load entry is unpinned
455 + without waiting for the step-count delay."""
456 + s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=5)
457 + _seed_cached(s, "a", n_blocks=2)
458 +
459 + req = _Request([_Feature("a")], request_id="req_y")
460 + s.update_state_after_alloc(req, 0)
461 + meta = s.build_connector_meta(scheduler_output=None)
462 + assert "a" in meta.loads
463 +
464 + # "a" is pinned. Try to evict — must fail.
465 + s.update_state_after_alloc(
466 + _Request([_Feature("b", length=2)], request_id="req_z"), 0
467 + )
468 + meta2 = s.build_connector_meta(scheduler_output=None)
469 + assert "b" not in meta2.saves
470 +
471 + # req_y finishes — unpin fires in build_connector_meta (fast-path).
472 + s.build_connector_meta(_SchedulerOutput(finished_req_ids={"req_y"}))
473 +
474 + # Next step: "a" is now evictable, so "b" can be allocated.
475 + s.update_state_after_alloc(
476 + _Request([_Feature("b", length=2)], request_id="req_w"), 0
477 + )
478 + meta3 = s.build_connector_meta(scheduler_output=None)
479 + assert "b" in meta3.saves
480 + s.shutdown()
tests/v1/ec_connector/unit/cpu/scheduler/test_step_tracker.py new
+162
@@ -0,0 +1,162 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for StepTracker."""
4 +
5 +from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.step_tracker import (
6 + StepTracker,
7 +)
8 +
9 +
10 +def test_basic_step_count_expiry_depth_1():
11 + """With max_concurrent_batches=1, entry expires after 1 step."""
12 + t = StepTracker(max_concurrent_batches=1)
13 + t.add("h1", "req_a")
14 +
15 + # Step 1: entry was added, committed to slot. Deque fills to maxlen=1,
16 + # but no expiry yet (the entry is in the slot just committed).
17 + result = t.step(set())
18 + assert result == []
19 +
20 + # Step 2: deque is full, oldest slot (containing h1) expires.
21 + result = t.step(set())
22 + assert result == ["h1"]
23 +
24 +
25 +def test_basic_step_count_expiry_depth_2():
26 + """With max_concurrent_batches=2, entry expires after 2 steps."""
27 + t = StepTracker(max_concurrent_batches=2)
28 + t.add("h1", "req_a")
29 +
30 + assert t.step(set()) == [] # step 1: committed
31 + assert t.step(set()) == [] # step 2: deque fills but h1 in newer slot
32 + assert t.step(set()) == ["h1"] # step 3: h1's slot is oldest, expires
33 +
34 +
35 +def test_first_finish_fast_path():
36 + """Entry is returned immediately when request_id appears in finished."""
37 + t = StepTracker(max_concurrent_batches=3)
38 + t.add("h1", "req_a")
39 +
40 + result = t.step({"req_a"})
41 + assert result == ["h1"]
42 +
43 +
44 +def test_no_double_processing_fast_path_then_expiry():
45 + """Entry fast-pathed should NOT fire again on deque expiry."""
46 + t = StepTracker(max_concurrent_batches=1)
47 + t.add("h1", "req_a")
48 +
49 + # Fast-path fires h1.
50 + result = t.step({"req_a"})
51 + assert result == ["h1"]
52 +
53 + # Normal expiry would fire now, but h1 is already processed.
54 + result = t.step(set())
55 + assert result == []
56 +
57 +
58 +def test_no_double_processing_expiry_then_fast_path():
59 + """Entry expired via step-count should NOT fire again via fast-path."""
60 + t = StepTracker(max_concurrent_batches=1)
61 + t.add("h1", "req_a")
62 +
63 + t.step(set()) # commit
64 + result = t.step(set()) # expiry fires h1
65 + assert result == ["h1"]
66 +
67 + # req_a finishes later — must not re-fire.
68 + result = t.step({"req_a"})
69 + assert result == []
70 +
71 +
72 +def test_multiplicity_same_hash_different_requests():
73 + """Two adds of the same mm_hash (different requests) → returned twice."""
74 + t = StepTracker(max_concurrent_batches=3)
75 + t.add("h1", "req_a")
76 + t.add("h1", "req_b")
77 +
78 + # Both requests finish in the same step.
79 + result = t.step({"req_a", "req_b"})
80 + assert sorted(result) == ["h1", "h1"]
81 +
82 +
83 +def test_mixed_expiry_and_fast_path():
84 + """Some entries expire, others fast-path, in the same step."""
85 + t = StepTracker(max_concurrent_batches=1)
86 + t.add("h1", "req_a")
87 +
88 + t.step(set()) # commit h1's slot
89 +
90 + # Add h2 in the next step.
91 + t.add("h2", "req_b")
92 +
93 + # This step: h1's slot expires AND req_b finishes.
94 + result = t.step({"req_b"})
95 + assert set(result) == {"h1", "h2"}
96 +
97 +
98 +def test_drain_all_returns_unprocessed():
99 + """drain_all returns everything pending and clears state."""
100 + t = StepTracker(max_concurrent_batches=3)
101 + t.add("h1", "req_a")
102 + t.add("h2", "req_b")
103 + t.step(set()) # commit both
104 + t.add("h3", "req_c") # in _current, not yet committed
105 +
106 + result = t.drain_all()
107 + assert set(result) == {"h1", "h2", "h3"}
108 +
109 + # After drain, everything is empty.
110 + assert t.step(set()) == []
111 + assert t.drain_all() == []
112 +
113 +
114 +def test_drain_all_skips_already_processed():
115 + """drain_all should not re-return entries already fast-pathed."""
116 + t = StepTracker(max_concurrent_batches=3)
117 + t.add("h1", "req_a")
118 + t.add("h2", "req_b")
119 +
120 + # Fast-path h1.
121 + result = t.step({"req_a"})
122 + assert result == ["h1"]
123 +
124 + # drain should only return h2.
125 + result = t.drain_all()
126 + assert result == ["h2"]
127 +
128 +
129 +def test_empty_step_returns_nothing():
130 + """step() with nothing pending returns empty list."""
131 + t = StepTracker(max_concurrent_batches=2)
132 + assert t.step(set()) == []
133 + assert t.step({"req_x"}) == []
134 +
135 +
136 +def test_add_after_step_goes_to_next_slot():
137 + """Entries added after a step go into the next slot, not the current."""
138 + t = StepTracker(max_concurrent_batches=1)
139 +
140 + t.step(set()) # empty step
141 + t.add("h1", "req_a")
142 +
143 + # h1 is in _current. Next step commits it.
144 + t.step(set())
145 +
146 + # Now the deque is full (1 slot with h1). Next step expires it.
147 + result = t.step(set())
148 + assert result == ["h1"]
149 +
150 +
151 +def test_fast_path_on_current_entries():
152 + """Entries still in _current (not yet committed) can be fast-pathed."""
153 + t = StepTracker(max_concurrent_batches=3)
154 + t.add("h1", "req_a")
155 +
156 + # req_a finishes in the same step h1 was added.
157 + result = t.step({"req_a"})
158 + assert result == ["h1"]
159 +
160 + # Entry is committed to slot but marked processed — won't re-fire.
161 + for _ in range(5):
162 + assert t.step(set()) == []
tests/v1/ec_connector/unit/cpu/test_connector.py new
+47
@@ -0,0 +1,47 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from unittest.mock import MagicMock
4 +
5 +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
6 +from vllm.distributed.ec_transfer.ec_connector.cpu.connector import ECCPUConnector
7 +from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
8 +
9 +
10 +def _cfg():
11 + ec = MagicMock()
12 + ec.is_ec_producer = True
13 + ec.is_ec_consumer = True
14 + cfg = MagicMock()
15 + cfg.ec_transfer_config = ec
16 + return cfg
17 +
18 +
19 +def test_scheduler_role_builds_only_scheduler(monkeypatch):
20 + fake_sched = MagicMock()
21 + monkeypatch.setattr(ECCPUConnector, "_make_scheduler", lambda self, cfg: fake_sched)
22 + c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER)
23 + assert c.connector_scheduler is fake_sched
24 + assert c.connector_worker is None
25 + c.has_cache_item("x")
26 + fake_sched.has_cache_item.assert_called_once_with("x")
27 +
28 +
29 +def test_worker_role_builds_only_worker(monkeypatch):
30 + fake_worker = MagicMock()
31 + monkeypatch.setattr(ECCPUConnector, "_make_worker", lambda self, cfg: fake_worker)
32 + c = ECCPUConnector(_cfg(), ECConnectorRole.WORKER)
33 + assert c.connector_worker is fake_worker
34 + assert c.connector_scheduler is None
35 +
36 +
37 +def test_request_finished_inherited_noop(monkeypatch):
38 + monkeypatch.setattr(
39 + ECCPUConnector, "_make_scheduler", lambda self, cfg: MagicMock()
40 + )
41 + c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER)
42 + assert c.request_finished(MagicMock()) == (False, None)
43 +
44 +
45 +def test_factory_registered():
46 + cls = ECConnectorFactory._registry["ECCPUConnector"]()
47 + assert cls is ECCPUConnector
tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py new
+170
@@ -0,0 +1,170 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for ECSharedRegion (the mmap substrate)."""
4 +
5 +import os
6 +import uuid
7 +from unittest.mock import MagicMock, patch
8 +
9 +import pytest
10 +import torch
11 +
12 +from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
13 + ECSharedRegion,
14 + _wait_for_file_size,
15 +)
16 +
17 +
18 +def _make_region(num_blocks: int = 8, block_size_bytes: int = 64) -> ECSharedRegion:
19 + return ECSharedRegion(
20 + engine_id=str(uuid.uuid4()),
21 + num_blocks=num_blocks,
22 + block_size_bytes=block_size_bytes,
23 + )
24 +
25 +
26 +@pytest.fixture
27 +def region() -> ECSharedRegion:
28 + r = _make_region()
29 + yield r
30 + r.cleanup()
31 +
32 +
33 +# ── mmap sharing between two instances ───────────────────────────────────────
34 +
35 +
36 +def test_second_instance_opens_existing_file_and_shares_memory():
37 + instance_id = str(uuid.uuid4())
38 + r1 = ECSharedRegion(
39 + engine_id=instance_id,
40 + num_blocks=4,
41 + block_size_bytes=64,
42 + )
43 + try:
44 + r2 = ECSharedRegion(
45 + engine_id=instance_id,
46 + num_blocks=4,
47 + block_size_bytes=64,
48 + )
49 + try:
50 + assert r1._is_creator
51 + assert not r2._is_creator
52 + # Both map the same physical pages; writes via r1 are visible via r2.
53 + r1.blocks[0, :4] = torch.tensor([10, 20, 30, 40], dtype=torch.int8)
54 + assert r2.blocks[0, :4].tolist() == [10, 20, 30, 40]
55 + finally:
56 + r2.cleanup()
57 + finally:
58 + r1.cleanup()
59 +
60 +
61 +def test_only_creator_unlinks_file_on_cleanup():
62 + """Critical contract: if the non-creator unlinks, the creator's mmap path
63 + becomes a dangling backing file and a third opener would create a new one
64 + out from under the creator."""
65 + instance_id = str(uuid.uuid4())
66 + r1 = ECSharedRegion(
67 + engine_id=instance_id,
68 + num_blocks=4,
69 + block_size_bytes=64,
70 + )
71 + path = r1._mmap_path
72 + r2 = ECSharedRegion(
73 + engine_id=instance_id,
74 + num_blocks=4,
75 + block_size_bytes=64,
76 + )
77 +
78 + # Non-creator goes away first — file must still be on disk for r1.
79 + r2.cleanup()
80 + assert os.path.exists(path), "non-creator cleanup must not unlink the file"
81 +
82 + # Creator goes away — file is removed.
83 + r1.cleanup()
84 + assert not os.path.exists(path), "creator cleanup must unlink the file"
85 +
86 +
87 +# ── _wait_for_file_size ──────────────────────────────────────────────────────
88 +
89 +
90 +def test_wait_for_file_size_returns_when_already_big_enough(tmp_path):
91 + """The fast path: file already at expected size — return immediately."""
92 + p = tmp_path / "f.bin"
93 + p.write_bytes(b"\x00" * 128)
94 + fd = os.open(str(p), os.O_RDONLY)
95 + try:
96 + _wait_for_file_size(fd, expected_size=128, timeout=1.0) # must not raise
97 + finally:
98 + os.close(fd)
99 +
100 +
101 +def test_wait_for_file_size_times_out_when_file_stays_empty(tmp_path):
102 + p = tmp_path / "f.bin"
103 + p.write_bytes(b"")
104 + fd = os.open(str(p), os.O_RDONLY)
105 + try:
106 + with pytest.raises(TimeoutError):
107 + _wait_for_file_size(fd, expected_size=4096, timeout=0.05)
108 + finally:
109 + os.close(fd)
110 +
111 +
112 +# ── pin_memory ───────────────────────────────────────────────────────────────
113 +
114 +
115 +def test_pin_memory_success_sets_flag(region):
116 + """When cudaHostRegister returns 0, _is_pinned flips to True
117 + and cleanup will correspondingly call cudaHostUnregister."""
118 + fake_cudart = MagicMock()
119 + success = MagicMock()
120 + success.value = 0
121 + fake_cudart.cudaHostRegister.return_value = success
122 + fake_cudart.cudaHostUnregister.return_value = success
123 +
124 + with (
125 + patch("torch.cuda.is_available", return_value=True),
126 + patch("torch.cuda.cudart", return_value=fake_cudart),
127 + ):
128 + region.pin_memory()
129 + assert region._is_pinned is True
130 + # cleanup must pair with cudaHostUnregister exactly once.
131 + region.cleanup()
132 + fake_cudart.cudaHostUnregister.assert_called_once()
133 +
134 +
135 +def test_pin_memory_failure_leaves_flag_false():
136 + """If cudaHostRegister fails (non-zero), don't pretend it succeeded —
137 + cleanup must NOT call cudaHostUnregister on memory we never registered."""
138 + r = _make_region()
139 + try:
140 + fake_cudart = MagicMock()
141 + fail = MagicMock()
142 + fail.value = 1 # non-zero == error
143 + fake_cudart.cudaHostRegister.return_value = fail
144 +
145 + with (
146 + patch("torch.cuda.is_available", return_value=True),
147 + patch("torch.cuda.cudart", return_value=fake_cudart),
148 + ):
149 + r.pin_memory()
150 + assert r._is_pinned is False
151 + # Now run cleanup and verify cudaHostUnregister was NOT called.
152 + r.cleanup()
153 + fake_cudart.cudaHostUnregister.assert_not_called()
154 + finally:
155 + r.cleanup()
156 +
157 +
158 +def test_pin_memory_noop_without_cuda(region):
159 + """pin_memory is a no-op when CUDA is not available."""
160 + with patch("torch.cuda.is_available", return_value=False):
161 + region.pin_memory()
162 + assert region._is_pinned is False
163 +
164 +
165 +# ── cleanup idempotency ───────────────────────────────────────────────────────
166 +
167 +
168 +def test_cleanup_is_idempotent(region):
169 + region.cleanup()
170 + region.cleanup() # fixture calls a third time — must not raise
tests/v1/ec_connector/unit/cpu/worker/test_worker.py new
+598
@@ -0,0 +1,598 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for ECCPUWorker.
4 +
5 +The byte-level tests exercise real CUDA stream/event coordination
6 +against a real ``ECSharedRegion`` mmap and are skipped on hosts without
7 +CUDA. The lifecycle tests don't fire any ``torch.cuda.*`` primitives and
8 +run anywhere.
9 +
10 +Mocking policy
11 +--------------
12 +- ``create_ec_shared_region`` is patched to inject a real ``ECSharedRegion``
13 + (mmap backed) with deterministic small dimensions, so byte assertions
14 + hit the actual mmap. The real responsibility of the worker — block
15 + index → mmap byte mapping, stream coordination, dtype/shape — is
16 + exercised end to end against a real GPU.
17 +- ``is_pin_memory_available`` is patched per-test so the rank-gating
18 + path can be probed deterministically without depending on host caps.
19 +- The shutdown test bypasses ``__init__`` via ``object.__new__`` and
20 + injects a ``MagicMock`` region.
21 +"""
22 +
23 +import contextlib
24 +import logging
25 +import uuid
26 +from unittest.mock import MagicMock, Mock, patch
27 +
28 +import pytest
29 +import torch
30 +
31 +from vllm.config import VllmConfig
32 +from vllm.config.parallel import ParallelConfig
33 +from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
34 + ECCPUConnectorMetadata,
35 +)
36 +from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
37 + ECSharedRegion,
38 +)
39 +from vllm.distributed.ec_transfer.ec_connector.cpu.worker import ECCPUWorker
40 +
41 +# ── shape constants ──────────────────────────────────────────────────────────
42 +
43 +# hidden_dim=8, dtype=fp16 (2 bytes) → block_size_bytes = 16, so one row of
44 +# encoder output fits in exactly one block.
45 +_HIDDEN_DIM = 8
46 +_DTYPE = torch.float16
47 +_BLOCK_SIZE_BYTES = _HIDDEN_DIM * _DTYPE.itemsize
48 +_NUM_BLOCKS = 8
49 +
50 +_requires_cuda = pytest.mark.skipif(
51 + not torch.cuda.is_available(),
52 + reason="exercises real CUDA stream/event coordination in ECCPUWorker",
53 +)
54 +
55 +
56 +def _make_region() -> ECSharedRegion:
57 + """Fresh region backed by a real per-test mmap file."""
58 + return ECSharedRegion(
59 + engine_id=str(uuid.uuid4()),
60 + num_blocks=_NUM_BLOCKS,
61 + block_size_bytes=_BLOCK_SIZE_BYTES,
62 + )
63 +
64 +
65 +def _vllm_config(rank: int = 0) -> Mock:
66 + cfg = Mock(spec=VllmConfig)
67 + cfg.parallel_config = Mock(spec=ParallelConfig)
68 + cfg.parallel_config.rank = rank
69 + cfg.model_config = Mock()
70 + cfg.model_config.dtype = _DTYPE
71 + return cfg
72 +
73 +
74 +def _meta(
75 + *, saves: dict | None = None, loads: dict | None = None
76 +) -> ECCPUConnectorMetadata:
77 + return ECCPUConnectorMetadata(saves=saves or {}, loads=loads or {})
78 +
79 +
80 +# ── fixtures ─────────────────────────────────────────────────────────────────
81 +
82 +
83 +@pytest.fixture
84 +def make_worker():
85 + """Factory that builds an ``ECCPUWorker`` over a real ``ECSharedRegion``.
86 +
87 + Tracks the workers and regions it creates and tears them down at the
88 + end of the test, so each test starts from a fresh /dev/shm file.
89 + """
90 + workers: list[ECCPUWorker] = []
91 + regions: list[ECSharedRegion] = []
92 +
93 + def factory(
94 + *,
95 + rank: int = 0,
96 + tp_rank: int = 0,
97 + pcp_rank: int = 0,
98 + pin_memory_available: bool = False,
99 + ) -> ECCPUWorker:
100 + region = _make_region()
101 + regions.append(region)
102 + pcp_group = Mock()
103 + pcp_group.rank_in_group = pcp_rank
104 + with (
105 + patch(
106 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.create_ec_shared_region",
107 + return_value=region,
108 + ),
109 + patch(
110 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.is_pin_memory_available",
111 + return_value=pin_memory_available,
112 + ),
113 + patch(
114 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.get_tensor_model_parallel_rank",
115 + return_value=tp_rank,
116 + ),
117 + patch(
118 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.get_pcp_group",
119 + return_value=pcp_group,
120 + ),
121 + ):
122 + worker = ECCPUWorker(_vllm_config(rank=rank))
123 + workers.append(worker)
124 + return worker
125 +
126 + yield factory
127 +
128 + for worker in workers:
129 + with contextlib.suppress(Exception):
130 + worker.shutdown()
131 + for region in regions:
132 + with contextlib.suppress(Exception):
133 + region.cleanup()
134 +
135 +
136 +# ── save_caches ──────────────────────────────────────────────────────────────
137 +
138 +
139 +@_requires_cuda
140 +@pytest.mark.parametrize(
141 + "n_elements,n_blocks",
142 + [
143 + (3 * _HIDDEN_DIM, 3), # exact fit: 3 full blocks, no remainder
144 + (3 * _HIDDEN_DIM + 4, 4), # partial last block: 4 fp16 trail in slot 3
145 + ],
146 + ids=["exact-fit", "partial-last-block"],
147 +)
148 +def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks):
149 + """``save_caches`` + ``flush_saves`` copies the source GPU tensor's bytes
150 + into the block IDs named by ``meta.saves``, in the order they appear,
151 + via a single batched GPU→CPU copy.
152 +
153 + When ``total_bytes`` is not a multiple of ``block_size_bytes`` the last
154 + block is partially written; the unwritten tail must remain whatever
155 + was there before. Blocks not named in ``meta.saves`` must be untouched.
156 + """
157 + worker = make_worker()
158 + sentinel = 0x5A
159 + worker._region.blocks.fill_(sentinel)
160 +
161 + src = torch.arange(n_elements, dtype=_DTYPE, device="cuda")
162 + expected_bytes = src.cpu().reshape(-1).view(torch.uint8)
163 + total_bytes = n_elements * _DTYPE.itemsize
164 +
165 + block_ids = [7, 2, 5, 0][:n_blocks]
166 + worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))
167 + worker.flush_saves()
168 +
169 + for slot, block_idx in enumerate(block_ids):
170 + block_byte_start = slot * _BLOCK_SIZE_BYTES
171 + block_byte_end = min(block_byte_start + _BLOCK_SIZE_BYTES, total_bytes)
172 + n_written = block_byte_end - block_byte_start
173 + actual_bytes = worker._region.blocks[block_idx, :n_written].view(torch.uint8)
174 + assert torch.equal(
175 + actual_bytes, expected_bytes[block_byte_start:block_byte_end]
176 + ), f"block {block_idx} (slot {slot}) bytes mismatch"
177 + if n_written < _BLOCK_SIZE_BYTES:
178 + tail = worker._region.blocks[block_idx, n_written:]
179 + assert torch.all(tail == sentinel), (
180 + f"block {block_idx} tail was overwritten"
181 + )
182 +
183 + for idx in set(range(_NUM_BLOCKS)) - set(block_ids):
184 + assert torch.all(worker._region.blocks[idx] == sentinel), (
185 + f"block {idx} (unassigned) was overwritten"
186 + )
187 +
188 +
189 +def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker):
190 + """When the scheduler hasn't pre-allocated blocks for ``mm_hash``,
191 + ``save_caches`` + ``flush_saves`` is a pure no-op."""
192 + worker = make_worker()
193 + sentinel = 0x42
194 + worker._region.blocks.fill_(sentinel)
195 +
196 + worker.save_caches({}, "h", _meta(saves={}))
197 + worker.flush_saves()
198 +
199 + assert torch.all(worker._region.blocks == sentinel)
200 +
201 +
202 +@pytest.mark.parametrize(
203 + "tp_rank,pcp_rank",
204 + [(1, 0), (0, 1), (1, 1)],
205 + ids=["tp1-pcp0", "tp0-pcp1", "tp1-pcp1"],
206 +)
207 +def test_save_caches_noop_for_non_save_rank(make_worker, tp_rank, pcp_rank):
208 + """Only TP rank 0 + PCP rank 0 writes to mmap. All other combinations
209 + must skip the save entirely."""
210 + worker = make_worker(tp_rank=tp_rank, pcp_rank=pcp_rank)
211 + sentinel = 0x42
212 + worker._region.blocks.fill_(sentinel)
213 +
214 + worker.save_caches(
215 + {"h": torch.zeros(8, dtype=_DTYPE)}, "h", _meta(saves={"h": [0]})
216 + )
217 + worker.flush_saves()
218 +
219 + assert torch.all(worker._region.blocks == sentinel)
220 +
221 +
222 +def test_save_caches_raises_when_allocated_blocks_too_small(make_worker):
223 + """``save_caches`` must raise ``AssertionError`` when the encoder output
224 + is larger than the allocated block space."""
225 + worker = make_worker()
226 + src = torch.zeros(3 * _HIDDEN_DIM, dtype=_DTYPE)
227 + with pytest.raises(AssertionError, match="exceeds allocated blocks"):
228 + worker.save_caches({"h": src}, "h", _meta(saves={"h": [0, 1]}))
229 +
230 +
231 +@_requires_cuda
232 +def test_save_caches_batches_multiple_hashes(make_worker):
233 + """Multiple save_caches calls are batched into a single flush."""
234 + worker = make_worker()
235 + sentinel = 0x5A
236 + worker._region.blocks.fill_(sentinel)
237 +
238 + src_a = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device="cuda")
239 + src_b = torch.arange(_HIDDEN_DIM, 2 * _HIDDEN_DIM, dtype=_DTYPE, device="cuda")
240 +
241 + cache = {"a": src_a, "b": src_b}
242 + worker.save_caches(cache, "a", _meta(saves={"a": [1], "b": [3]}))
243 + worker.save_caches(cache, "b", _meta(saves={"a": [1], "b": [3]}))
244 + worker.flush_saves()
245 +
246 + expected_a = src_a.cpu().view(torch.uint8)
247 + expected_b = src_b.cpu().view(torch.uint8)
248 + actual_a = worker._region.blocks[1].view(torch.uint8)
249 + actual_b = worker._region.blocks[3].view(torch.uint8)
250 + assert torch.equal(actual_a, expected_a)
251 + assert torch.equal(actual_b, expected_b)
252 +
253 + # Unassigned blocks untouched.
254 + for idx in (0, 2, 4, 5, 6, 7):
255 + assert torch.all(worker._region.blocks[idx] == sentinel)
256 +
257 +
258 +# ── start_load_caches ────────────────────────────────────────────────────────
259 +
260 +
261 +@_requires_cuda
262 +def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker):
263 + """Single batched load across all hashes with correct byte→dtype→shape."""
264 + worker = make_worker()
265 + n_blocks = 3
266 + src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape(
267 + n_blocks, _HIDDEN_DIM
268 + )
269 + src_int8 = src_orig.view(torch.int8).reshape(n_blocks, _BLOCK_SIZE_BYTES)
270 +
271 + block_ids = [3, 1, 6]
272 + for i, idx in enumerate(block_ids):
273 + worker._region.blocks[idx].copy_(src_int8[i])
274 +
275 + encoder_cache: dict[str, torch.Tensor] = {}
276 + worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids}))
277 +
278 + out = encoder_cache["h"]
279 + assert out.is_cuda, "consumer worker must place the tensor on the GPU"
280 + assert out.shape == (n_blocks, _HIDDEN_DIM)
281 + assert out.dtype == _DTYPE
282 + assert torch.equal(out.cpu(), src_orig)
283 +
284 +
285 +@_requires_cuda
286 +def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker):
287 + """If ``encoder_cache`` already holds the ``mm_hash``, the worker must
288 + not overwrite it."""
289 + worker = make_worker()
290 + worker._region.blocks[0].fill_(0x42)
291 +
292 + sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device="cuda")
293 + encoder_cache = {"h": sentinel}
294 + worker.start_load_caches(encoder_cache, _meta(loads={"h": [0]}))
295 +
296 + assert encoder_cache["h"] is sentinel, (
297 + "existing encoder_cache entry must not be replaced"
298 + )
299 +
300 +
301 +@_requires_cuda
302 +def test_start_load_caches_noop_when_loads_is_empty(make_worker):
303 + """When ``meta.loads`` is empty the early-return must fire."""
304 + worker = make_worker()
305 + encoder_cache: dict[str, torch.Tensor] = {}
306 + worker.start_load_caches(encoder_cache, _meta(loads={}))
307 +
308 + assert encoder_cache == {}
309 +
310 +
311 +@_requires_cuda
312 +def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):
313 + """Cached entries are preserved while new ones are loaded."""
314 + worker = make_worker()
315 + n_blocks = 2
316 + src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape(
317 + n_blocks, _HIDDEN_DIM
318 + )
319 + src_int8 = src_orig.view(torch.int8).reshape(n_blocks, _BLOCK_SIZE_BYTES)
320 + new_block_ids = [4, 2]
321 + for i, idx in enumerate(new_block_ids):
322 + worker._region.blocks[idx].copy_(src_int8[i])
323 +
324 + cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device="cuda")
325 + encoder_cache: dict[str, torch.Tensor] = {"cached_h": cached_tensor}
326 + worker.start_load_caches(
327 + encoder_cache,
328 + _meta(loads={"cached_h": [0], "new_h": new_block_ids}),
329 + )
330 +
331 + assert encoder_cache["cached_h"] is cached_tensor
332 + assert "new_h" in encoder_cache
333 + out = encoder_cache["new_h"]
334 + assert out.shape == (n_blocks, _HIDDEN_DIM)
335 + assert out.dtype == _DTYPE
336 + assert torch.equal(out.cpu(), src_orig)
337 +
338 +
339 +@_requires_cuda
340 +@pytest.mark.parametrize(
341 + "tp_rank,pcp_rank",
342 + [(0, 0), (1, 0), (0, 1), (1, 1)],
343 + ids=["tp0-pcp0", "tp1-pcp0", "tp0-pcp1", "tp1-pcp1"],
344 +)
345 +def test_start_load_caches_works_on_all_ranks(make_worker, tp_rank, pcp_rank):
346 + """All TP/PCP ranks must load from mmap — loads are NOT gated like saves."""
347 + worker = make_worker(tp_rank=tp_rank, pcp_rank=pcp_rank)
348 + n_blocks = 2
349 + src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape(
350 + n_blocks, _HIDDEN_DIM
351 + )
352 + src_int8 = src_orig.view(torch.int8).reshape(n_blocks, _BLOCK_SIZE_BYTES)
353 + block_ids = [1, 3]
354 + for i, idx in enumerate(block_ids):
355 + worker._region.blocks[idx].copy_(src_int8[i])
356 +
357 + encoder_cache: dict[str, torch.Tensor] = {}
358 + worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids}))
359 +
360 + out = encoder_cache["h"]
361 + assert out.is_cuda
362 + assert torch.equal(out.cpu(), src_orig)
363 +
364 +
365 +# ── round-trip ───────────────────────────────────────────────────────────────
366 +
367 +
368 +@_requires_cuda
369 +def test_save_then_load_round_trips_bytes(make_worker):
370 + """Full producer→mmap→consumer byte path in one shot."""
371 + worker = make_worker()
372 + n_blocks = 3
373 + block_ids = [5, 1, 6]
374 +
375 + src = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device="cuda").reshape(
376 + n_blocks, _HIDDEN_DIM
377 + )
378 + worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))
379 + worker.flush_saves()
380 +
381 + encoder_cache: dict[str, torch.Tensor] = {}
382 + worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids}))
383 +
384 + out = encoder_cache["h"]
385 + assert out.shape == src.shape
386 + assert out.dtype == src.dtype
387 + assert torch.equal(out.cpu(), src.cpu())
388 +
389 +
390 +# ── buffer recycling ────────────────────────────────────────────────────────
391 +
392 +
393 +@_requires_cuda
394 +def test_buffer_pool_is_reused_across_save_steps(make_worker):
395 + """After flush_saves, descriptor buffers are returned to the pool and
396 + reused on the next flush — no reallocation."""
397 + worker = make_worker()
398 + src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device="cuda")
399 +
400 + worker.save_caches({"h": src}, "h", _meta(saves={"h": [0]}))
401 + worker.flush_saves()
402 +
403 + assert len(worker._buf_pool._pool) == 1
404 + buf_id = id(worker._buf_pool._pool[0].src_ptrs)
405 +
406 + # Second step reuses the same buffer.
407 + worker.save_caches({"h": src}, "h", _meta(saves={"h": [1]}))
408 + worker.flush_saves()
409 +
410 + assert len(worker._buf_pool._pool) == 1
411 + assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id
412 +
413 +
414 +@_requires_cuda
415 +def test_buffer_pool_is_reused_across_load_steps(make_worker):
416 + """After start_load_caches, descriptor buffers are returned to the pool
417 + and reused on the next call."""
418 + worker = make_worker()
419 + worker._region.blocks[0].fill_(0x01)
420 + worker._region.blocks[1].fill_(0x02)
421 +
422 + encoder_cache: dict[str, torch.Tensor] = {}
423 + worker.start_load_caches(encoder_cache, _meta(loads={"a": [0]}))
424 +
425 + assert len(worker._buf_pool._pool) == 1
426 + buf_id = id(worker._buf_pool._pool[0].src_ptrs)
427 +
428 + encoder_cache2: dict[str, torch.Tensor] = {}
429 + worker.start_load_caches(encoder_cache2, _meta(loads={"b": [1]}))
430 +
431 + assert len(worker._buf_pool._pool) == 1
432 + assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id
433 +
434 +
435 +# ── stream management ────────────────────────────────────────────────────────
436 +
437 +
438 +@_requires_cuda
439 +def test_stream_initialized_at_construction(make_worker):
440 + """``_load_stream`` must be a fully initialized CUDA stream."""
441 + worker = make_worker()
442 + assert isinstance(worker._load_stream, torch.cuda.Stream)
443 +
444 +
445 +# ── lifecycle ────────────────────────────────────────────────────────────────
446 +
447 +
448 +@pytest.mark.parametrize(
449 + "pin_available,expected_pinned",
450 + [
451 + (True, True),
452 + (False, False),
453 + ],
454 + ids=["available", "unavailable"],
455 +)
456 +def test_init_pins_memory_when_available(pin_available, expected_pinned):
457 + """Every worker must call ``pin_memory()`` when the platform allows it."""
458 + region = _make_region()
459 + region.pin_memory = MagicMock() # type: ignore[method-assign]
460 +
461 + pcp_group = Mock()
462 + pcp_group.rank_in_group = 0
463 + with (
464 + patch(
465 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.create_ec_shared_region",
466 + return_value=region,
467 + ),
468 + patch(
469 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.is_pin_memory_available",
470 + return_value=pin_available,
471 + ),
472 + patch(
473 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.get_tensor_model_parallel_rank",
474 + return_value=0,
475 + ),
476 + patch(
477 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.get_pcp_group",
478 + return_value=pcp_group,
479 + ),
480 + ):
481 + ECCPUWorker(_vllm_config(rank=0))
482 +
483 + try:
484 + if expected_pinned:
485 + region.pin_memory.assert_called_once()
486 + else:
487 + region.pin_memory.assert_not_called()
488 + finally:
489 + region.cleanup()
490 +
491 +
492 +def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):
493 + """``shutdown`` must always call ``region.cleanup`` — and must never
494 + raise."""
495 + worker = object.__new__(ECCPUWorker)
496 + mock_region = Mock(spec=ECSharedRegion)
497 + worker._region = mock_region
498 + worker._load_stream = MagicMock()
499 + worker._save_bufs = None
500 + worker._save_count = 0
501 +
502 + worker.shutdown()
503 + worker._load_stream.synchronize.assert_called_once()
504 + mock_region.cleanup.assert_called_once()
505 +
506 + mock_region.cleanup.side_effect = RuntimeError("boom")
507 + with caplog_vllm.at_level(logging.DEBUG, logger="vllm"):
508 + worker.shutdown() # exception must be swallowed
509 +
510 + assert mock_region.cleanup.call_count == 2
511 + assert any(
512 + "worker region cleanup failed" in r.message
513 + for r in caplog_vllm.records
514 + if r.levelno == logging.DEBUG
515 + )
516 +
517 +
518 +# ── e2e: scheduler + worker pipeline ────────────────────────────────────────
519 +
520 +
521 +@_requires_cuda
522 +def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):
523 + """Full pipeline: scheduler allocates blocks, worker saves GPU tensor to
524 + mmap via flush_saves, scheduler marks ready after step delay, worker
525 + loads from mmap back to GPU, and the result matches the original.
526 +
527 + Exercises the real scheduler + worker cooperation through a shared
528 + ECSharedRegion, with real CUDA transfers and stream coordination.
529 + """
530 + import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
531 + from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import (
532 + ECCPUScheduler,
533 + )
534 +
535 + # Build worker first (gets us a real region).
536 + worker = make_worker()
537 + region = worker._region
538 +
539 + # Build scheduler sharing the same region.
540 + monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)
541 +
542 + class _EC:
543 + is_ec_producer = True
544 + is_ec_consumer = True
545 +
546 + class _Cfg:
547 + ec_transfer_config = _EC()
548 + max_concurrent_batches = 1
549 +
550 + scheduler = ECCPUScheduler(_Cfg())
551 +
552 + # -- Step 1: scheduler allocates, worker saves --
553 + n_blocks = 3
554 + src = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device="cuda").reshape(
555 + n_blocks, _HIDDEN_DIM
556 + )
557 +
558 + class _Pos:
559 + offset = 0
560 + length = n_blocks
561 +
562 + class _Feature:
563 + identifier = "img_001"
564 + mm_position = _Pos()
565 +
566 + class _Request:
567 + request_id = "req_e2e"
568 + mm_features = [_Feature()]
569 +
570 + scheduler.update_state_after_alloc(_Request(), 0)
571 + meta_save = scheduler.build_connector_meta(scheduler_output=None)
572 + assert "img_001" in meta_save.saves
573 +
574 + encoder_cache = {"img_001": src}
575 + worker.save_caches(encoder_cache, "img_001", meta_save)
576 + worker.flush_saves()
577 +
578 + # -- Step 2: readiness delay (max_concurrent_batches=1) --
579 + meta_step2 = scheduler.build_connector_meta(scheduler_output=None)
580 + assert scheduler.has_cache_item("img_001") is True
581 + assert meta_step2.loads == {} # no load requested yet
582 +
583 + # -- Step 3: scheduler emits load, worker loads --
584 + scheduler.update_state_after_alloc(_Request(), 0)
585 + meta_load = scheduler.build_connector_meta(scheduler_output=None)
586 + assert "img_001" in meta_load.loads
587 + assert meta_load.loads["img_001"] == meta_save.saves["img_001"]
588 +
589 + load_cache: dict[str, torch.Tensor] = {}
590 + worker.start_load_caches(load_cache, meta_load)
591 +
592 + out = load_cache["img_001"]
593 + assert out.is_cuda
594 + assert out.shape == src.shape
595 + assert out.dtype == src.dtype
596 + assert torch.equal(out.cpu(), src.cpu())
597 +
598 + scheduler.shutdown()
tests/v1/ec_connector/unit/test_ec_cpu_connector.py new
+158
@@ -0,0 +1,158 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +E2E test for the ECCPUConnector.
5 +
6 +Verifies:
7 +- Accuracy: outputs from EC CPU cache match fresh encoder computation.
8 +- Latency: loading from EC CPU cache is faster than a cold encoder run.
9 +
10 +Requires a CUDA GPU and the Qwen2-VL-2B-Instruct model.
11 +"""
12 +
13 +import time
14 +
15 +import pytest
16 +from tqdm import tqdm
17 +
18 +from vllm import LLM, SamplingParams
19 +from vllm.assets.image import ImageAsset
20 +from vllm.config import ECTransferConfig
21 +from vllm.platforms import current_platform
22 +
23 +MODEL = "Qwen/Qwen2-VL-2B-Instruct"
24 +EC_CPU_BYTES = 500 << 20 # 500 MiB
25 +
26 +
27 +def _build_image_prompt(image):
28 + """Build a Qwen2-VL prompt with a single image."""
29 + return {
30 + "prompt": (
31 + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
32 + "<|im_start|>user\n"
33 + "<|vision_start|><|image_pad|><|vision_end|>"
34 + "Describe this image briefly.<|im_end|>\n"
35 + "<|im_start|>assistant\n"
36 + ),
37 + "multi_modal_data": {"image": image},
38 + }
39 +
40 +
41 +def _make_unique_image(base_image, iteration: int):
42 + """Create a unique image by modifying a pixel to change the mm_hash."""
43 + img = base_image.copy()
44 + img.putpixel((0, 0), (iteration % 256, iteration // 256 % 256, 0))
45 + return img
46 +
47 +
48 +def _wait_for_ec_ready(llm: LLM) -> None:
49 + """Send a dummy request to force a scheduler step.
50 +
51 + After generate() returns, the EC entry's data is in the mmap but
52 + readiness may not have fired yet (first-finish fires in the NEXT
53 + scheduler step). This forces that step.
54 + """
55 + llm.generate("hi", SamplingParams(max_tokens=1), use_tqdm=False)
56 +
57 +
58 +def _latency_test(llm: LLM) -> None:
59 + """Verify EC CPU cache hit is faster than cold encoder computation."""
60 + if not current_platform.is_cuda():
61 + pytest.skip("Latency test requires CUDA")
62 +
63 + sampling_params = SamplingParams(max_tokens=1, temperature=0)
64 + base_image = ImageAsset("stop_sign").pil_image
65 +
66 + num_times_ec_better_than_cold = 0
67 + num_tests = 10
68 + total_cold_time = 0.0
69 + total_ec_hit_time = 0.0
70 +
71 + for i in tqdm(range(num_tests), desc="EC latency test"):
72 + # Use a unique image each iteration so the cold run is truly cold
73 + # (EC has never seen this mm_hash before).
74 + image = _make_unique_image(base_image, i)
75 + prompt = _build_image_prompt(image)
76 +
77 + # Cold run: encoder computes from scratch, saves to EC CPU.
78 + llm.llm_engine.reset_encoder_cache()
79 + start_time = time.time()
80 + llm.generate([prompt], sampling_params, use_tqdm=False)
81 + cold_time = time.time() - start_time
82 + total_cold_time += cold_time
83 +
84 + # Ensure EC entry is marked ready.
85 + _wait_for_ec_ready(llm)
86 +
87 + # Clear GPU encoder cache but keep EC CPU cache intact.
88 + llm.llm_engine.reset_encoder_cache()
89 +
90 + # EC hit: scheduler loads from CPU mmap instead of re-encoding.
91 + start_time = time.time()
92 + llm.generate([prompt], sampling_params, use_tqdm=False)
93 + ec_hit_time = time.time() - start_time
94 + total_ec_hit_time += ec_hit_time
95 +
96 + if ec_hit_time < cold_time:
97 + num_times_ec_better_than_cold += 1
98 +
99 + print("Average times:")
100 + print(f" Cold: {total_cold_time * 1000 / num_tests:.2f}ms")
101 + print(f" EC hit: {total_ec_hit_time * 1000 / num_tests:.2f}ms")
102 +
103 + assert num_times_ec_better_than_cold >= 0.8 * num_tests, (
104 + f"EC hit was faster only {num_times_ec_better_than_cold}/{num_tests} "
105 + f"times (expected >= 80%)"
106 + )
107 +
108 +
109 +def _accuracy_test(llm: LLM) -> None:
110 + """Verify EC CPU cache produces correct outputs."""
111 + sampling_params = SamplingParams(max_tokens=8, temperature=0)
112 + image = ImageAsset("stop_sign").pil_image
113 + prompt = _build_image_prompt(image)
114 +
115 + # Seed: generate with fresh encoder to establish baseline output.
116 + llm.llm_engine.reset_encoder_cache()
117 + baseline = llm.generate([prompt], sampling_params, use_tqdm=False)
118 + baseline_text = baseline[0].outputs[0].text
119 +
120 + # Ensure EC entry is marked ready.
121 + _wait_for_ec_ready(llm)
122 +
123 + # Clear GPU encoder cache so subsequent runs go through EC CPU path.
124 + llm.llm_engine.reset_encoder_cache()
125 +
126 + # Generate multiple times from EC cache and check correctness.
127 + test_count = 20
128 + results = llm.generate([prompt] * test_count, sampling_params, use_tqdm=False)
129 +
130 + success_count = sum(1 for r in results if r.outputs[0].text == baseline_text)
131 + assert success_count >= 0.5 * test_count, (
132 + f"Only {success_count}/{test_count} outputs matched baseline "
133 + f"(expected >= 50%). Baseline: {baseline_text!r}"
134 + )
135 +
136 +
137 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
138 +def test_ec_cpu_offloading() -> None:
139 + """Tests ECCPUConnector accuracy and latency with a VLM model."""
140 + ec_transfer_config = ECTransferConfig(
141 + ec_connector="ECCPUConnector",
142 + ec_role="ec_both",
143 + ec_connector_extra_config={"ec_cpu_bytes": EC_CPU_BYTES},
144 + )
145 +
146 + llm = LLM(
147 + model=MODEL,
148 + max_model_len=2048,
149 + gpu_memory_utilization=0.5,
150 + enforce_eager=True,
151 + ec_transfer_config=ec_transfer_config,
152 + )
153 +
154 + try:
155 + _latency_test(llm)
156 + _accuracy_test(llm)
157 + finally:
158 + del llm
tests/v1/ec_connector/unit/test_metadata.py new
+96
@@ -0,0 +1,96 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for ECCPUConnectorMetadata."""
4 +
5 +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
6 +from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
7 + ECCPUConnectorMetadata,
8 +)
9 +
10 +
11 +def test_metadata_is_ec_connector_metadata():
12 + """ECCPUConnectorMetadata must be a subclass of ECConnectorMetadata."""
13 + meta = ECCPUConnectorMetadata()
14 + assert isinstance(meta, ECConnectorMetadata)
15 +
16 +
17 +def test_metadata_initialization():
18 + """ECCPUConnectorMetadata initializes with empty saves and loads dicts."""
19 + meta = ECCPUConnectorMetadata()
20 + assert hasattr(meta, "saves")
21 + assert hasattr(meta, "loads")
22 + assert meta.saves == {}
23 + assert meta.loads == {}
24 +
25 +
26 +def test_metadata_saves_dict_operations():
27 + """The saves dict supports standard dict operations."""
28 + meta = ECCPUConnectorMetadata()
29 +
30 + # Add multiple entries
31 + meta.saves["mm_hash_1"] = [0, 1, 2]
32 + meta.saves["mm_hash_2"] = [10, 11]
33 +
34 + assert "mm_hash_1" in meta.saves
35 + assert "mm_hash_2" in meta.saves
36 + assert len(meta.saves) == 2
37 + assert meta.saves["mm_hash_1"] == [0, 1, 2]
38 + assert meta.saves["mm_hash_2"] == [10, 11]
39 +
40 +
41 +def test_metadata_loads_dict_operations():
42 + """The loads dict supports standard dict operations."""
43 + meta = ECCPUConnectorMetadata()
44 +
45 + # Add multiple entries
46 + meta.loads["mm_hash_1"] = [5, 6, 7]
47 + meta.loads["mm_hash_2"] = [100]
48 +
49 + assert "mm_hash_1" in meta.loads
50 + assert "mm_hash_2" in meta.loads
51 + assert len(meta.loads) == 2
52 + assert meta.loads["mm_hash_1"] == [5, 6, 7]
53 + assert meta.loads["mm_hash_2"] == [100]
54 +
55 +
56 +def test_metadata_saves_and_loads_are_independent():
57 + """saves and loads dicts are independent."""
58 + meta = ECCPUConnectorMetadata()
59 +
60 + meta.saves["key"] = [1, 2]
61 + meta.loads["key"] = [3, 4]
62 +
63 + # Same key but different dicts should not interfere
64 + assert meta.saves["key"] == [1, 2]
65 + assert meta.loads["key"] == [3, 4]
66 +
67 + # Modifying one should not affect the other
68 + meta.saves["other"] = [99]
69 + assert "other" not in meta.loads
70 +
71 +
72 +def test_metadata_block_indices_are_integers():
73 + """Block indices in saves/loads lists should be integers."""
74 + meta = ECCPUConnectorMetadata()
75 +
76 + meta.saves["hash"] = [0, 100, 999, 65535]
77 + meta.loads["hash"] = [1, 2, 3]
78 +
79 + for idx in meta.saves["hash"]:
80 + assert isinstance(idx, int)
81 + for idx in meta.loads["hash"]:
82 + assert isinstance(idx, int)
83 +
84 +
85 +def test_metadata_multiple_instances_are_independent():
86 + """Different ECCPUConnectorMetadata instances should not share state."""
87 + meta1 = ECCPUConnectorMetadata()
88 + meta2 = ECCPUConnectorMetadata()
89 +
90 + meta1.saves["x"] = [1]
91 + meta2.saves["y"] = [2]
92 +
93 + assert "x" in meta1.saves
94 + assert "x" not in meta2.saves
95 + assert "y" not in meta1.saves
96 + assert "y" in meta2.saves
vllm/distributed/ec_transfer/ec_connector/cpu/__init__.py
vllm/distributed/ec_transfer/ec_connector/cpu/common.py new
+89
@@ -0,0 +1,89 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared types for the ECCPUConnector scheduler and worker delegates."""
4 +
5 +from dataclasses import dataclass, field
6 +from typing import TYPE_CHECKING
7 +
8 +import torch
9 +
10 +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
11 +from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
12 + ECSharedRegion,
13 +)
14 +
15 +if TYPE_CHECKING:
16 + from vllm.config import VllmConfig
17 +
18 +
19 +@dataclass
20 +class ECCPUConnectorMetadata(ECConnectorMetadata):
21 + """Per-step scheduler → worker payload for the ECCPUConnector.
22 +
23 + Populated by `ECCPUScheduler.build_connector_meta`; consumed by
24 + `ECCPUWorker` via the mixin's `bind_connector_metadata`.
25 + """
26 +
27 + # Producer role: mm_hashes the scheduler has just allocated CPU
28 + # blocks for this step; the worker's save_caches copies
29 + # encoder_cache[mm_hash] → mmap at these block IDs.
30 + saves: dict[str, list[int]] = field(default_factory=dict)
31 +
32 + # Consumer role: mm_hashes whose bytes are available in the local mmap;
33 + # the worker's start_load_caches copies mmap[block_ids] → GPU
34 + # encoder_cache.
35 + loads: dict[str, list[int]] = field(default_factory=dict)
36 +
37 +
38 +def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int:
39 + """Return the per-token hidden dimension for encoder cache entries.
40 +
41 + For most models this equals the LLM's hidden size. Qwen3-VL (and any
42 + future model with deepstack visual encoding) is an exception: the ViT
43 + concatenates its own output with features from N decoder layers before
44 + storing in encoder_cache, producing a tensor of width
45 + ``out_hidden_size * (1 + N)`` per visual token. Using the plain LLM
46 + hidden size would under-allocate EC blocks and silently truncate the
47 + transferred data, leading to a shape mismatch on the consumer.
48 + """
49 + model_config = vllm_config.model_config
50 + hf_config = getattr(model_config, "hf_config", None)
51 + vision_config = (
52 + getattr(hf_config, "vision_config", None) if hf_config is not None else None
53 + )
54 + if vision_config is not None:
55 + out_hidden_size = getattr(vision_config, "out_hidden_size", None)
56 + deepstack_indexes = getattr(vision_config, "deepstack_visual_indexes", None)
57 + if out_hidden_size is not None and deepstack_indexes:
58 + return out_hidden_size * (1 + len(deepstack_indexes))
59 + return model_config.get_inputs_embeds_size()
60 +
61 +
62 +def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion:
63 + """Build the EC mmap region from `vllm_config`.
64 +
65 + Both `ECCPUScheduler` and `ECCPUWorker` call this to get the same
66 + shared region (same engine_id, same block_size_bytes).
67 + """
68 + ec_config = vllm_config.ec_transfer_config
69 + assert ec_config is not None, "ec_transfer_config required to build region"
70 +
71 + dp_rank = vllm_config.parallel_config.data_parallel_rank
72 + engine_id = f"{vllm_config.instance_id}_dp{dp_rank}"
73 +
74 + dtype = vllm_config.model_config.dtype
75 + hidden_dim = _get_encoder_cache_hidden_dim(vllm_config)
76 + element_size = torch.empty(0, dtype=dtype).element_size()
77 + block_size_bytes = hidden_dim * element_size
78 +
79 + cpu_bytes = ec_config.ec_connector_extra_config.get("ec_cpu_bytes")
80 + if not cpu_bytes:
81 + raise ValueError("ec_cpu_bytes must be specified in ec_connector_extra_config")
82 + cpu_bytes = int(cpu_bytes)
83 + num_blocks = cpu_bytes // block_size_bytes
84 +
85 + return ECSharedRegion(
86 + engine_id=engine_id,
87 + num_blocks=num_blocks,
88 + block_size_bytes=block_size_bytes,
89 + )
vllm/distributed/ec_transfer/ec_connector/cpu/connector.py new
+121
@@ -0,0 +1,121 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ECCPUConnector — CPU encoder-cache offloading.
4 +
5 +A role-routed shell: one instance per process. The scheduler delegate owns the
6 +mmap region and offload bookkeeping; the worker delegate owns the mmap view and
7 +GPU copy plumbing. An ec_both instance reuses encoder outputs it has already
8 +offloaded to CPU instead of recomputing them.
9 +"""
10 +
11 +from typing import TYPE_CHECKING
12 +
13 +import torch
14 +
15 +from vllm.distributed.ec_transfer.ec_connector.base import (
16 + ECConnectorBase,
17 + ECConnectorRole,
18 +)
19 +from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
20 + ECCPUConnectorMetadata,
21 +)
22 +from vllm.logger import init_logger
23 +
24 +if TYPE_CHECKING:
25 + from vllm.config import VllmConfig
26 + from vllm.v1.core.sched.output import SchedulerOutput
27 + from vllm.v1.request import Request
28 +
29 +logger = init_logger(__name__)
30 +
31 +
32 +class ECCPUConnector(ECConnectorBase):
33 + """EC connector that offloads encoder cache to a shared CPU mmap region."""
34 +
35 + def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole) -> None:
36 + super().__init__(vllm_config=vllm_config, role=role)
37 +
38 + self.connector_worker = None
39 + self.connector_scheduler = None
40 +
41 + if role == ECConnectorRole.WORKER:
42 + self.connector_worker = self._make_worker(vllm_config)
43 + elif role == ECConnectorRole.SCHEDULER:
44 + self.connector_scheduler = self._make_scheduler(vllm_config)
45 + else:
46 + raise ValueError(f"Unknown ECConnectorRole: {role}")
47 +
48 + # Construction seams.
49 + def _make_worker(self, vllm_config: "VllmConfig"):
50 + # Deferred import: the worker module touches torch/CUDA at import time
51 + # via the region, so keep that cost off the scheduler path.
52 + from vllm.distributed.ec_transfer.ec_connector.cpu.worker import (
53 + ECCPUWorker,
54 + )
55 +
56 + return ECCPUWorker(vllm_config)
57 +
58 + def _make_scheduler(self, vllm_config: "VllmConfig"):
59 + from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import (
60 + ECCPUScheduler,
61 + )
62 +
63 + return ECCPUScheduler(vllm_config)
64 +
65 + # Worker-side forwarders.
66 + def start_load_caches(
67 + self, encoder_cache: dict[str, torch.Tensor], **kwargs
68 + ) -> None:
69 + assert self.connector_worker is not None
70 + metadata = self._get_connector_metadata()
71 + assert isinstance(metadata, ECCPUConnectorMetadata)
72 + self.connector_worker.start_load_caches(
73 + encoder_cache, connector_metadata=metadata
74 + )
75 +
76 + def save_caches(
77 + self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs
78 + ) -> None:
79 + assert self.connector_worker is not None
80 + metadata = self._get_connector_metadata()
81 + assert isinstance(metadata, ECCPUConnectorMetadata)
82 + self.connector_worker.save_caches(
83 + encoder_cache, mm_hash, connector_metadata=metadata
84 + )
85 +
86 + # Scheduler-side forwarders.
87 + def has_cache_item(self, identifier: str) -> bool:
88 + assert self.connector_scheduler is not None
89 + return self.connector_scheduler.has_cache_item(identifier)
90 +
91 + def ensure_cache_available(
92 + self, request: "Request", num_computed_tokens: int
93 + ) -> bool:
94 + assert self.connector_scheduler is not None
95 + return self.connector_scheduler.ensure_cache_available(
96 + request, num_computed_tokens
97 + )
98 +
99 + def update_state_after_alloc(self, request: "Request", index: int) -> None:
100 + assert self.connector_scheduler is not None
101 + self.connector_scheduler.update_state_after_alloc(request, index)
102 +
103 + def build_connector_meta(
104 + self, scheduler_output: "SchedulerOutput"
105 + ) -> ECCPUConnectorMetadata:
106 + assert self.connector_scheduler is not None
107 + return self.connector_scheduler.build_connector_meta(scheduler_output)
108 +
109 + def get_finished(
110 + self, finished_req_ids: set[str]
111 + ) -> tuple[set[str] | None, set[str] | None]:
112 + if self.connector_worker is not None:
113 + self.connector_worker.flush_saves()
114 + return None, None
115 +
116 + # Shared.
117 + def shutdown(self) -> None:
118 + if self.connector_scheduler is not None:
119 + self.connector_scheduler.shutdown()
120 + if self.connector_worker is not None:
121 + self.connector_worker.shutdown()
vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py new
+159
@@ -0,0 +1,159 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Lightweight mmap-backed shared memory region for encoder cache (EC) data.
5 +
6 +Modeled after SharedOffloadRegion (vllm/v1/kv_offload/cpu/) but simplified
7 +for EC: flat shared layout, no multi-tensor cursor, no block_size_factor.
8 +"""
9 +
10 +import mmap
11 +import os
12 +import time
13 +
14 +import torch
15 +
16 +from vllm.logger import init_logger
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0):
22 + """Spin-wait until the file reaches expected_size (creator truncated it)."""
23 + deadline = time.monotonic() + timeout
24 + while True:
25 + if os.fstat(fd).st_size >= expected_size:
26 + return
27 + if time.monotonic() > deadline:
28 + raise TimeoutError(
29 + f"Timed out waiting for EC mmap file to reach {expected_size} bytes"
30 + )
31 + time.sleep(0.005)
32 +
33 +
34 +class ECSharedRegion:
35 + """Flat mmap-backed memory region shared across TP workers for
36 + encoder cache blocks.
37 +
38 + Layout: (num_blocks, block_size_bytes) — contiguous, no per-worker
39 + interleaving. All workers map the same file and see identical data.
40 +
41 + File path: /dev/shm/vllm_ec_{engine_id}.mmap
42 +
43 + This class owns only the shared memory substrate (mmap lifecycle, the
44 + `blocks` view, CUDA host registration). Block allocation and eviction
45 + are tracked by `EmbeddingCache` in the scheduler process.
46 + """
47 +
48 + def __init__(
49 + self,
50 + engine_id: str,
51 + num_blocks: int,
52 + block_size_bytes: int,
53 + ) -> None:
54 + self.num_blocks = num_blocks
55 + self.block_size_bytes = block_size_bytes
56 +
57 + total_size_bytes = num_blocks * block_size_bytes
58 + # Path in /dev/shm (tmpfs); unique per engine instance.
59 + self._mmap_path = f"/dev/shm/vllm_ec_{engine_id}.mmap"
60 + # True for the process that created the file (responsible for unlink).
61 + self._is_creator = False
62 + # True after successful cudaHostRegister (cleanup must unregister).
63 + self._is_pinned = False
64 +
65 + # File descriptor for the shared memory backing file.
66 + try:
67 + self._fd: int | None = os.open(
68 + self._mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600
69 + )
70 + os.ftruncate(self._fd, total_size_bytes)
71 + self._is_creator = True
72 + logger.info(
73 + "Created EC mmap file %s (%.2f MiB)",
74 + self._mmap_path,
75 + total_size_bytes / (1 << 20),
76 + )
77 + except FileExistsError:
78 + self._fd = os.open(self._mmap_path, os.O_RDWR)
79 + try:
80 + _wait_for_file_size(self._fd, total_size_bytes)
81 + except Exception:
82 + os.close(self._fd)
83 + self._fd = None
84 + raise
85 + logger.info("Opened existing EC mmap file %s", self._mmap_path)
86 +
87 + # MAP_SHARED mmap over _fd; all processes see the same pages.
88 + self._mmap_obj: mmap.mmap | None = mmap.mmap(
89 + self._fd,
90 + total_size_bytes,
91 + flags=mmap.MAP_SHARED,
92 + prot=mmap.PROT_READ | mmap.PROT_WRITE,
93 + )
94 +
95 + if self._is_creator:
96 + _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
97 + self._mmap_obj.madvise(_MADV_POPULATE_WRITE, 0, total_size_bytes)
98 +
99 + # (num_blocks, block_size_bytes) int8 tensor over the mmap buffer.
100 + self.blocks: torch.Tensor = torch.frombuffer(
101 + memoryview(self._mmap_obj), dtype=torch.int8
102 + ).view(num_blocks, block_size_bytes)
103 + # Cached for cudaHostRegister/Unregister and pointer math.
104 + self._blocks_ptr: int = self.blocks.data_ptr()
105 + self._blocks_nbytes: int = self.blocks.nbytes
106 +
107 + def pin_memory(self) -> None:
108 + """Register the entire mmap as CUDA pinned memory for fast DMA.
109 +
110 + Each TP worker process owns its own virtual address for the same
111 + shared physical pages, so every process must register independently.
112 + No-op when CUDA is not available or already pinned.
113 + """
114 + if self._is_pinned or not torch.cuda.is_available():
115 + return
116 + result = torch.cuda.cudart().cudaHostRegister(
117 + self._blocks_ptr, self._blocks_nbytes, 0
118 + )
119 + if result.value != 0:
120 + logger.warning(
121 + "cudaHostRegister failed (code=%d) — "
122 + "transfers will still work but may be slower (unpinned DMA)",
123 + result.value,
124 + )
125 + else:
126 + logger.debug("cudaHostRegister %.2f MB", self._blocks_nbytes / 1e6)
127 + self._is_pinned = True
128 +
129 + def cleanup(self) -> None:
130 + """Tear down the region. Lifecycle method; no concurrent access."""
131 + logger.info("Starting ECSharedRegion cleanup...")
132 + if self._is_creator:
133 + try:
134 + os.unlink(self._mmap_path)
135 + logger.info("Removed EC mmap file %s", self._mmap_path)
136 + except Exception:
137 + logger.warning(
138 + "Failed to unlink path %s", self._mmap_path, exc_info=True
139 + )
140 + self._is_creator = False
141 + if self._is_pinned:
142 + result = torch.cuda.cudart().cudaHostUnregister(self._blocks_ptr)
143 + if result.value != 0:
144 + logger.warning("cudaHostUnregister failed (code=%d)", result)
145 + self._is_pinned = False
146 + if hasattr(self, "blocks"):
147 + del self.blocks
148 + if self._mmap_obj:
149 + try:
150 + self._mmap_obj.close()
151 + except Exception:
152 + logger.warning("Failed to close mmap_obj", exc_info=True)
153 + self._mmap_obj = None
154 + if self._fd is not None:
155 + try:
156 + os.close(self._fd)
157 + except Exception:
158 + logger.warning("Failed to close fd %s", self._fd, exc_info=True)
159 + self._fd = None
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py new
+122
@@ -0,0 +1,122 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ECCPUScheduler — CPU offload scheduler delegate.
4 +
5 +Owns the mmap region and the embedding cache, and handles the producer
6 +(GPU->CPU offload) and consumer (CPU->GPU reload) scheduler-side logic
7 +for the ECCPUConnector.
8 +"""
9 +
10 +from typing import TYPE_CHECKING
11 +
12 +from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
13 + ECCPUConnectorMetadata,
14 + create_ec_shared_region,
15 +)
16 +from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
17 + EmbeddingCache,
18 +)
19 +from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.step_tracker import (
20 + StepTracker,
21 +)
22 +from vllm.logger import init_logger
23 +
24 +if TYPE_CHECKING:
25 + from vllm.config import VllmConfig
26 + from vllm.v1.core.sched.output import SchedulerOutput
27 + from vllm.v1.request import Request
28 +
29 +logger = init_logger(__name__)
30 +
31 +
32 +class ECCPUScheduler:
33 + """Scheduler delegate for the ECCPUConnector."""
34 +
35 + def __init__(self, vllm_config: "VllmConfig") -> None:
36 + ec_config = vllm_config.ec_transfer_config
37 + assert ec_config is not None
38 + self._is_producer: bool = ec_config.is_ec_producer
39 + self._is_consumer: bool = ec_config.is_ec_consumer
40 +
41 + self._region = create_ec_shared_region(vllm_config)
42 + # Block allocator + LRU eviction policy for the shared region.
43 + self._cache = EmbeddingCache(self._region.num_blocks)
44 +
45 + max_batches = vllm_config.max_concurrent_batches
46 + # Delays mark_ready until the GPU→mmap DMA is guaranteed complete.
47 + self._ready_tracker = StepTracker(max_batches)
48 + # Delays unpin until the mmap→GPU DMA is guaranteed complete.
49 + self._unpin_tracker = StepTracker(max_batches)
50 +
51 + # mm_hash → block IDs allocated this step for GPU→mmap saves.
52 + self._pending_saves: dict[str, list[int]] = {}
53 + # mm_hash → block IDs to load from mmap→GPU this step.
54 + self._pending_loads: dict[str, list[int]] = {}
55 +
56 + def has_cache_item(self, identifier: str) -> bool:
57 + if not self._is_consumer:
58 + return False
59 + entry = self._cache.get(identifier)
60 + return entry is not None and entry.ready
61 +
62 + def ensure_cache_available(
63 + self, request: "Request", num_computed_tokens: int
64 + ) -> bool:
65 + return True
66 +
67 + def update_state_after_alloc(self, request: "Request", index: int) -> None:
68 + feature = request.mm_features[index]
69 + mm_hash = feature.identifier
70 +
71 + if self._is_producer and self._cache.get(mm_hash) is None:
72 + entry = self._cache.alloc(mm_hash, feature.mm_position.length)
73 + if entry is not None:
74 + self._pending_saves[mm_hash] = list(entry.block_ids)
75 + self._ready_tracker.add(mm_hash, request.request_id)
76 +
77 + if self._is_consumer and mm_hash not in self._pending_loads:
78 + entry = self._cache.get(mm_hash)
79 + if entry is not None and entry.ready:
80 + self._cache.pin(mm_hash)
81 + self._pending_loads[mm_hash] = list(entry.block_ids)
82 + self._unpin_tracker.add(mm_hash, request.request_id)
83 +
84 + def build_connector_meta(
85 + self, scheduler_output: "SchedulerOutput"
86 + ) -> ECCPUConnectorMetadata:
87 + finished = (
88 + scheduler_output.finished_req_ids if scheduler_output is not None else set()
89 + )
90 +
91 + for key in self._ready_tracker.step(finished):
92 + entry = self._cache.get(key)
93 + if entry is not None and not entry.ready:
94 + self._cache.mark_ready(key)
95 +
96 + for key in self._unpin_tracker.step(finished):
97 + self._cache.unpin(key)
98 +
99 + meta = ECCPUConnectorMetadata()
100 + if self._is_producer:
101 + meta.saves = self._pending_saves
102 + self._pending_saves = {}
103 + if self._is_consumer:
104 + meta.loads = self._pending_loads
105 + self._pending_loads = {}
106 + return meta
107 +
108 + def shutdown(self) -> None:
109 + # drain_all() covers both entries still in _current (never
110 + # consumed by build_connector_meta) and entries in slots.
111 + self._pending_loads.clear()
112 + for mm_hash in self._unpin_tracker.drain_all():
113 + self._cache.unpin(mm_hash)
114 + self._ready_tracker.drain_all()
115 +
116 + self._is_producer = False
117 + self._is_consumer = False
118 +
119 + try:
120 + self._region.cleanup()
121 + except Exception:
122 + logger.debug("ec: region cleanup failed", exc_info=True)
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py new
+131
@@ -0,0 +1,131 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""EmbeddingCache — named-entry block cache with FIFO eviction.
4 +
5 +Manages a fixed pool of block IDs keyed by content identity (mm_hash).
6 +Entries transition through: not-ready → ready (evictable) → pinned.
7 +Eviction targets ready + unpinned entries in FIFO order.
8 +
9 +All public methods are thread-safe.
10 +"""
11 +
12 +import threading
13 +from collections import OrderedDict
14 +
15 +
16 +class CacheEntry:
17 + """A single cache entry. Read `.block_ids` and `.ready` freely;
18 + mutations only through EmbeddingCache methods."""
19 +
20 + __slots__ = ("block_ids", "_pin_count")
21 +
22 + def __init__(self, block_ids: tuple[int, ...]) -> None:
23 + self.block_ids = block_ids
24 + self._pin_count = -1 # not ready
25 +
26 + @property
27 + def ready(self) -> bool:
28 + return self._pin_count >= 0
29 +
30 + @property
31 + def evictable(self) -> bool:
32 + return self._pin_count == 0
33 +
34 + def mark_ready(self):
35 + self._pin_count = 0
36 +
37 + def pin(self):
38 + self._pin_count += 1
39 +
40 +
41 +class EmbeddingCache:
42 + """Fixed-size block cache with FIFO eviction.
43 +
44 + Entries are keyed by content identity (e.g. mm_hash). Blocks are
45 + allocated from a free-set; when space is needed, ready + unpinned
46 + entries are evicted oldest-first.
47 +
48 + The caller is responsible for deciding *when* to call ``mark_ready``
49 + (e.g. after enough engine steps have elapsed for the worker to have
50 + completed the write).
51 + """
52 +
53 + def __init__(self, num_blocks: int) -> None:
54 + self._num_blocks = num_blocks
55 + self._lock = threading.Lock()
56 + # Available block IDs, LIFO stack.
57 + self._free_blocks: list[int] = list(range(num_blocks))
58 + # Key → live entry (all states: not-ready, ready, pinned).
59 + self._entries: dict[str, CacheEntry] = {}
60 + # Ready + unpinned entries in insertion order (FIFO eviction).
61 + self._entries_free_list: OrderedDict[str, None] = OrderedDict()
62 + # Total blocks reclaimable via eviction (fail-fast check).
63 + self._evictable_block_count: int = 0
64 +
65 + def get(self, key: str) -> CacheEntry | None:
66 + """Return the entry for *key*, or None if not present."""
67 + with self._lock:
68 + return self._entries.get(key)
69 +
70 + def alloc(self, key: str, n_blocks: int) -> CacheEntry | None:
71 + """Allocate *n_blocks* for *key*, evicting as needed.
72 +
73 + The entry starts not-ready. Returns None if there is not enough
74 + space even after evicting all evictable entries.
75 + """
76 + with self._lock:
77 + assert key not in self._entries, (
78 + f"EmbeddingCache: duplicate alloc for {key!r}"
79 + )
80 + assert n_blocks <= self._num_blocks, (
81 + f"EmbeddingCache: {n_blocks} blocks requested but capacity "
82 + f"is {self._num_blocks}"
83 + )
84 + if len(self._free_blocks) + self._evictable_block_count < n_blocks:
85 + return None
86 + self._evict_until(n_blocks)
87 + block_ids = tuple(self._free_blocks.pop() for _ in range(n_blocks))
88 + entry = CacheEntry(block_ids)
89 + self._entries[key] = entry
90 + return entry
91 +
92 + def mark_ready(self, key: str) -> None:
93 + """Mark an entry as ready (data is CPU-visible)."""
94 + with self._lock:
95 + entry = self._entries[key]
96 + assert entry._pin_count == -1, (
97 + f"EmbeddingCache: mark_ready on already-ready entry {key!r}"
98 + )
99 + entry.mark_ready()
100 + self._entries_free_list[key] = None
101 + self._evictable_block_count += len(entry.block_ids)
102 +
103 + def pin(self, key: str) -> None:
104 + """Pin an entry (prevent eviction)."""
105 + with self._lock:
106 + entry = self._entries[key]
107 + assert entry.ready, f"EmbeddingCache: pin of not-ready entry {key!r}"
108 + if entry.evictable:
109 + del self._entries_free_list[key]
110 + self._evictable_block_count -= len(entry.block_ids)
111 + entry.pin()
112 +
113 + def unpin(self, key: str) -> None:
114 + """Unpin an entry. Asserts currently pinned."""
115 + with self._lock:
116 + entry = self._entries[key]
117 + assert entry._pin_count > 0, (
118 + f"EmbeddingCache: unpin of unpinned entry {key!r}"
119 + )
120 + entry._pin_count -= 1
121 + if entry._pin_count == 0:
122 + self._entries_free_list[key] = None
123 + self._evictable_block_count += len(entry.block_ids)
124 +
125 + def _evict_until(self, n_blocks: int) -> None:
126 + """Evict ready+unpinned entries FIFO until enough space. Lock held."""
127 + while len(self._free_blocks) < n_blocks:
128 + key, _ = self._entries_free_list.popitem(last=False)
129 + entry = self._entries.pop(key)
130 + self._evictable_block_count -= len(entry.block_ids)
131 + self._free_blocks.extend(entry.block_ids)
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/step_tracker.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""StepTracker — step-counting delay with first-finish fast-path.
4 +
5 +Manages a collection of (mm_hash, request_id) entries that need to be
6 +processed after a configurable number of engine steps have elapsed.
7 +Additionally, if the originating request finishes before the step-count
8 +expires, the entry is processed immediately (the GPU transfer is
9 +guaranteed complete by request finish time).
10 +
11 +Used by ECCPUScheduler: one instance for readiness marking, one for
12 +unpins.
13 +"""
14 +
15 +from collections import deque
16 +from dataclasses import dataclass
17 +
18 +
19 +@dataclass(slots=True)
20 +class _PendingEntry:
21 + mm_hash: str
22 + request_id: str
23 + processed: bool = False
24 +
25 +
26 +class StepTracker:
27 + """Track pending operations with step-count delay and first-finish.
28 +
29 + Each ``add(mm_hash, request_id)`` registers an operation that will be
30 + returned from ``step()`` either when:
31 + (a) ``max_concurrent_batches`` steps elapse (deque expiry), OR
32 + (b) ``request_id`` appears in ``finished_req_ids`` (first-finish).
33 +
34 + Each add is processed exactly once — never both (a) and (b).
35 + """
36 +
37 + def __init__(self, max_concurrent_batches: int) -> None:
38 + # One slot per past step; oldest at the right. When the deque is
39 + # full the rightmost slot is popped and its entries expire.
40 + self._slots: deque[list[_PendingEntry]] = deque(maxlen=max_concurrent_batches)
41 + # Reverse index: request_id → entries across all slots + _current.
42 + # Enables O(1) lookup for the first-finish fast-path.
43 + self._req_index: dict[str, list[_PendingEntry]] = {}
44 + # Entries added this step (via add()), committed to a new slot at
45 + # the start of the next step() call.
46 + self._current: list[_PendingEntry] = []
47 +
48 + def add(self, mm_hash: str, request_id: str) -> None:
49 + """Register an mm_hash needing processing after a delay."""
50 + entry = _PendingEntry(mm_hash=mm_hash, request_id=request_id)
51 + self._current.append(entry)
52 + self._req_index.setdefault(request_id, []).append(entry)
53 +
54 + def step(self, finished_req_ids: set[str]) -> list[str]:
55 + """Advance one step. Return mm_hashes ready for processing.
56 +
57 + Returns:
58 + List of mm_hashes to process. May contain duplicates if the
59 + same mm_hash was added multiple times (from different requests).
60 + """
61 + result: list[str] = []
62 +
63 + # Phase 1: first-finish fast-path.
64 + for req_id in finished_req_ids:
65 + entries = self._req_index.pop(req_id, None)
66 + if entries is None:
67 + continue
68 + for entry in entries:
69 + if not entry.processed:
70 + entry.processed = True
71 + result.append(entry.mm_hash)
72 +
73 + # Phase 2: step-count expiry.
74 + if len(self._slots) == self._slots.maxlen:
75 + expired_slot = self._slots.pop()
76 + for entry in expired_slot:
77 + if not entry.processed:
78 + entry.processed = True
79 + result.append(entry.mm_hash)
80 + self._cleanup_expired(expired_slot)
81 +
82 + # Phase 3: commit current entries to new slot.
83 + self._slots.appendleft(self._current)
84 + self._current = []
85 +
86 + return result
87 +
88 + def drain_all(self) -> list[str]:
89 + """Return all unprocessed entries and clear internal state."""
90 + result: list[str] = []
91 + for entry in self._current:
92 + if not entry.processed:
93 + entry.processed = True
94 + result.append(entry.mm_hash)
95 + self._current = []
96 +
97 + for slot in self._slots:
98 + for entry in slot:
99 + if not entry.processed:
100 + entry.processed = True
101 + result.append(entry.mm_hash)
102 + self._slots.clear()
103 + self._req_index.clear()
104 + return result
105 +
106 + def _cleanup_expired(self, slot: list[_PendingEntry]) -> None:
107 + """Prune reverse index for entries in an expired slot."""
108 + for entry in slot:
109 + entries_for_req = self._req_index.get(entry.request_id)
110 + if entries_for_req is None:
111 + # Already removed by fast-path (pop in step()).
112 + continue
113 + entries_for_req.remove(entry)
114 + if not entries_for_req:
115 + del self._req_index[entry.request_id]
vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py new
+202
@@ -0,0 +1,202 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Worker-side of the ECCPUConnector.
4 +
5 +Thin, stateless across steps: opens the shared mmap region and uses the
6 +per-step connector metadata (`ECCPUConnectorMetadata`) to decide which
7 +blocks to copy in each direction.
8 +"""
9 +
10 +from typing import TYPE_CHECKING
11 +
12 +import torch
13 +
14 +from vllm._custom_ops import swap_blocks_batch
15 +from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
16 + ECCPUConnectorMetadata,
17 + create_ec_shared_region,
18 +)
19 +from vllm.distributed.ec_transfer.ec_connector.cpu.worker.descriptor_buffers import (
20 + DescriptorBufferPool,
21 + DescriptorBuffers,
22 +)
23 +from vllm.distributed.parallel_state import (
24 + get_pcp_group,
25 + get_tensor_model_parallel_rank,
26 +)
27 +from vllm.logger import init_logger
28 +from vllm.platforms import current_platform
29 +from vllm.utils.platform_utils import is_pin_memory_available
30 +
31 +if TYPE_CHECKING:
32 + from vllm.config import VllmConfig
33 +
34 +logger = init_logger(__name__)
35 +
36 +
37 +class ECCPUWorker:
38 + """Worker-side delegate for the ECCPUConnector.
39 +
40 + - Producer role: copies `encoder_cache[mm_hash]` → `mmap[block_ids]`
41 + for each entry in `metadata.saves`. Descriptor buffers are filled
42 + directly in `save_caches`; the actual DMA is issued as a single
43 + batched call in `flush_saves`.
44 + - Consumer role: copies `mmap[block_ids]` → `encoder_cache[mm_hash]`
45 + for all entries in `metadata.loads` via a single `swap_blocks_batch`
46 + call on the load stream.
47 + - On `ec_both` nodes both paths run back-to-back in a single step.
48 + """
49 +
50 + def __init__(self, vllm_config: "VllmConfig") -> None:
51 + self._region = create_ec_shared_region(vllm_config)
52 + # Model dtype; used to reinterpret raw int8 blocks on load.
53 + self._dtype = vllm_config.model_config.dtype
54 +
55 + if is_pin_memory_available():
56 + self._region.pin_memory()
57 +
58 + # All TP/PCP ranks hold identical encoder output. Only one rank
59 + # per mmap needs to write — saves host memory bandwidth.
60 + # DCP is a subdivision of TP, so tp_rank==0 covers it.
61 + self._is_save_rank = (
62 + get_tensor_model_parallel_rank() == 0 and get_pcp_group().rank_in_group == 0
63 + )
64 +
65 + # Dedicated stream for async mmap→GPU loads (overlaps with compute).
66 + self._load_stream = current_platform.Stream()
67 +
68 + # Descriptor buffer pool (recycled across steps, shared by both paths).
69 + self._buf_pool = DescriptorBufferPool()
70 +
71 + # Active save buffer being filled during save_caches calls this step.
72 + self._save_bufs: DescriptorBuffers | None = None
73 + self._save_count: int = 0
74 +
75 + def save_caches(
76 + self,
77 + encoder_cache: dict[str, torch.Tensor],
78 + mm_hash: str,
79 + connector_metadata: ECCPUConnectorMetadata,
80 + ) -> None:
81 + """Fill descriptor buffers directly for batched flush."""
82 + if not self._is_save_rank:
83 + return
84 + block_ids = connector_metadata.saves.get(mm_hash)
85 + if block_ids is None:
86 + return
87 +
88 + src = encoder_cache[mm_hash]
89 + total_bytes = src.numel() * src.element_size()
90 + block_size = self._region.block_size_bytes
91 + allocated_bytes = len(block_ids) * block_size
92 + assert total_bytes <= allocated_bytes, (
93 + f"EC: encoder output exceeds allocated blocks for mm_hash={mm_hash}: "
94 + f"{total_bytes} bytes but only {allocated_bytes} allocated "
95 + f"({len(block_ids)} blocks × {block_size}). "
96 + f"shape={list(src.shape)}"
97 + )
98 +
99 + if self._save_bufs is None:
100 + total = sum(len(v) for v in connector_metadata.saves.values())
101 + self._save_bufs = self._buf_pool.acquire(total)
102 +
103 + assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel()
104 +
105 + src_ptrs, dst_ptrs, sizes = self._save_bufs
106 + src_base = src.view(-1).view(torch.uint8).data_ptr()
107 + dst_base = self._region.blocks.data_ptr()
108 + idx = self._save_count
109 +
110 + for i, block_idx in enumerate(block_ids):
111 + start = i * block_size
112 + src_ptrs[idx] = src_base + start
113 + dst_ptrs[idx] = dst_base + block_idx * block_size
114 + sizes[idx] = min(block_size, total_bytes - start)
115 + idx += 1
116 +
117 + self._save_count = idx
118 +
119 + def flush_saves(self) -> None:
120 + """Flush all accumulated saves in a single swap_blocks_batch call."""
121 + if self._save_count == 0:
122 + return
123 +
124 + bufs = self._save_bufs
125 + assert bufs is not None
126 + src_ptrs, dst_ptrs, sizes = bufs
127 + n = self._save_count
128 + swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n])
129 +
130 + self._buf_pool.release(bufs)
131 + self._save_bufs = None
132 + self._save_count = 0
133 +
134 + def start_load_caches(
135 + self,
136 + encoder_cache: dict[str, torch.Tensor],
137 + connector_metadata: ECCPUConnectorMetadata,
138 + ) -> None:
139 + """Consumer path: single batched copy of all loads from mmap→GPU."""
140 + if not connector_metadata.loads:
141 + return
142 +
143 + block_size = self._region.block_size_bytes
144 + blocks = self._region.blocks
145 + dtype = self._dtype
146 + device_type = current_platform.device_type
147 + src_base = blocks.data_ptr()
148 +
149 + # Pre-filter: only hashes not already in encoder_cache.
150 + load_items = {
151 + h: idxs
152 + for h, idxs in connector_metadata.loads.items()
153 + if h not in encoder_cache
154 + }
155 + if not load_items:
156 + return
157 +
158 + total_blocks = sum(len(idxs) for idxs in load_items.values())
159 +
160 + with current_platform.stream(self._load_stream):
161 + # Single contiguous destination buffer for all loads.
162 + dst_buf = torch.empty(
163 + total_blocks, block_size, dtype=torch.int8, device=device_type
164 + )
165 + dst_buf_base = dst_buf.data_ptr()
166 +
167 + bufs = self._buf_pool.acquire(total_blocks)
168 + src_ptrs = bufs.src_ptrs[:total_blocks]
169 + dst_ptrs = bufs.dst_ptrs[:total_blocks]
170 + sizes = bufs.sizes[:total_blocks]
171 + sizes[:] = block_size
172 +
173 + op_idx = 0
174 + for block_ids in load_items.values():
175 + for block_idx in block_ids:
176 + src_ptrs[op_idx] = src_base + block_idx * block_size
177 + dst_ptrs[op_idx] = dst_buf_base + op_idx * block_size
178 + op_idx += 1
179 +
180 + swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any=True)
181 +
182 + self._buf_pool.release(bufs)
183 +
184 + # Slice contiguous buffer into per-hash views.
185 + offset = 0
186 + for mm_hash, block_ids in load_items.items():
187 + n = len(block_ids)
188 + encoder_cache[mm_hash] = (
189 + dst_buf[offset : offset + n].view(dtype).reshape(n, -1)
190 + )
191 + offset += n
192 +
193 + current_platform.current_stream().wait_stream(self._load_stream)
194 +
195 + def shutdown(self) -> None:
196 + self._load_stream.synchronize()
197 + self._save_bufs = None
198 + self._save_count = 0
199 + try:
200 + self._region.cleanup()
201 + except Exception:
202 + logger.debug("EC: worker region cleanup failed", exc_info=True)
vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py new
+46
@@ -0,0 +1,46 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Reusable pool of (src_ptrs, dst_ptrs, sizes) int64 tensor triples.
4 +
5 +Used by ECCPUWorker to batch swap_blocks_batch descriptors without
6 +per-step allocation overhead.
7 +"""
8 +
9 +from typing import NamedTuple
10 +
11 +import torch
12 +
13 +
14 +class DescriptorBuffers(NamedTuple):
15 + src_ptrs: torch.Tensor
16 + dst_ptrs: torch.Tensor
17 + sizes: torch.Tensor
18 +
19 +
20 +class DescriptorBufferPool:
21 + """Pool of descriptor buffer triples for swap_blocks_batch.
22 +
23 + Each buffer is a `DescriptorBuffers` namedtuple of three 1-D int64
24 + tensors of equal length. Buffers are recycled across steps; if a
25 + returned buffer is too small it is discarded and a fresh one allocated.
26 + """
27 +
28 + def __init__(self) -> None:
29 + # LIFO stack of idle buffer triples.
30 + self._pool: list[DescriptorBuffers] = []
31 +
32 + def acquire(self, n: int) -> DescriptorBuffers:
33 + """Get a buffer triple with capacity >= *n*."""
34 + if self._pool:
35 + bufs = self._pool.pop()
36 + if bufs.src_ptrs.numel() >= n:
37 + return bufs
38 + return DescriptorBuffers(
39 + torch.empty(n, dtype=torch.int64),
40 + torch.empty(n, dtype=torch.int64),
41 + torch.empty(n, dtype=torch.int64),
42 + )
43 +
44 + def release(self, bufs: DescriptorBuffers) -> None:
45 + """Return a buffer triple to the pool for reuse."""
46 + self._pool.append(bufs)
vllm/distributed/ec_transfer/ec_connector/factory.py
+6
@@ -83,3 +83,9 @@ ECConnectorFactory.register_connector(
83 "vllm.distributed.ec_transfer.ec_connector.example_connector",
84 "ECExampleConnector",
85 )
86 +
87 +ECConnectorFactory.register_connector(
88 + "ECCPUConnector",
89 + "vllm.distributed.ec_transfer.ec_connector.cpu.connector",
90 + "ECCPUConnector",
91 +)