@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f8d38fbc87

[EC Connector] EC Offloading Connector use events instead of StepTracker (#49994)

Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

omerpaz95 committed Aug 27, 2026 at 06:41 UTC f8d38fbc879c3a223a633d4ec532bdde348570ec

15 files changed +1049 -645

tests/v1/ec_connector/__init__.py

tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py

+30 -1

@@ -21,13 +21,15 @@ def _cache(num_blocks: int = 8) -> EmbeddingCache:

21 # ── alloc ────────────────────────────────────────────────────────────────────

24 -def test_alloc_returns_unique_block_ids():

24 +def test_alloc_returns_unique_ascending_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 + # Ascending order lets consumers detect which blocks are adjacent.

32 + assert list(entry.block_ids) == sorted(entry.block_ids)

35 def test_alloc_returns_tuple():

@@ -161,6 +163,33 @@ def test_unpin_asserts_on_unpinned():

163 cache.unpin("a")

164

165

166 +# ── has_held_entries ─────────────────────────────────────────────────────────

167 +

168 +

169 +def test_has_held_entries_empty_cache():

170 + cache = _cache()

171 + assert cache.has_held_entries() is False

172 +

173 +

174 +def test_has_held_entries_true_while_not_ready():

175 + cache = _cache()

176 + cache.alloc("a", 2) # not ready → held

177 + assert cache.has_held_entries() is True

178 + cache.mark_ready("a") # ready + unpinned → settled

179 + assert cache.has_held_entries() is False

180 +

181 +

182 +def test_has_held_entries_true_while_pinned():

183 + cache = _cache()

184 + cache.alloc("a", 2)

185 + cache.mark_ready("a")

186 + assert cache.has_held_entries() is False

187 + cache.pin("a") # pinned → held

188 + assert cache.has_held_entries() is True

189 + cache.unpin("a")

190 + assert cache.has_held_entries() is False

191 +

192 +

193 # ── eviction ─────────────────────────────────────────────────────────────────

194

195

tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py

+160 -178

@@ -3,6 +3,11 @@

3 import uuid

5 import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod

6 +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config

7 +from vllm.config.ec_transfer import ECRole

8 +from vllm.distributed.ec_transfer.ec_connector.cpu.common import (

9 + ECCPUWorkerMetadata,

10 +)

11 from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (

12 ECSharedRegion,

13 )

@@ -37,13 +42,26 @@ class _Request:

42 self.request_id = request_id

45 +class _WorkerOutput:

46 + """Stand-in for ECConnectorOutput carrying a completion report.

47 +

48 + `saves` are mm_hashes; `loads` are transfer ids. One entry per reporting

49 + rank, so a list holding the same id twice is two ranks reporting it.

50 + """

51 +

52 + def __init__(self, *, saves=None, loads=None):

53 + self.ec_connector_worker_meta = ECCPUWorkerMetadata(

54 + completed_saves=list(saves or []),

55 + completed_loads=list(loads or []),

56 + )

57 +

58 +

59 def _make_scheduler(

60 monkeypatch,

61 num_blocks=_N,

43 - max_concurrent_batches=1,

62 *,

45 - is_producer=True,

46 - is_consumer=True,

63 + ec_role: ECRole = "ec_both",

64 + tensor_parallel_size=1,

65 ) -> ECCPUScheduler:

66 region = ECSharedRegion(

67 engine_id=str(uuid.uuid4()),

@@ -52,27 +70,22 @@ def _make_scheduler(

70 )

71 monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)

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

73 + return ECCPUScheduler(

74 + create_ec_vllm_config(

75 + ec_role=ec_role,

76 + tensor_parallel_size=tensor_parallel_size,

77 + )

78 + )

66 - @property

67 - def max_concurrent_batches(self):

68 - return _mcb

70 - return ECCPUScheduler(_Cfg())

81 +def _load_ids(meta) -> list[int]:

82 + """Transfer ids the scheduler dispatched in this step's metadata."""

83 + return [transfer_id for transfer_id, _ in meta.loads.values()]

73 -class _SchedulerOutput:

74 - def __init__(self, finished_req_ids=None):

75 - self.finished_req_ids = finished_req_ids or set()

86 +def _load_blocks(meta, mm_hash: str) -> list[int]:

87 + """Block ids the scheduler dispatched for one load."""

88 + return meta.loads[mm_hash][1]

91 def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int):

@@ -83,7 +96,9 @@ def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int):

97 def test_offload_reuse_cycle(monkeypatch):

98 s = _make_scheduler(monkeypatch)

86 - req = _Request([_Feature("h1", length=1)])

99 + # Two blocks, so the reload also pins down that every allocated block id

100 + # comes back in meta.loads, not just the first.

101 + req = _Request([_Feature("h1", length=2)])

102

103 # Step A: first sight — allocate and save.

104 assert s.has_cache_item("h1") is False

@@ -92,24 +107,24 @@ def test_offload_reuse_cycle(monkeypatch):

107 meta_a = s.build_connector_meta(scheduler_output=None)

108 assert "h1" in meta_a.saves

109 assert meta_a.loads == {}

95 - # Not ready yet (max_concurrent_batches=1 means ready next step).

110 + # Not ready until the worker reports the save memcpy complete.

111 assert s.has_cache_item("h1") is False

112

98 - # Step B: _step_readiness marks h1 ready (no new allocs).

99 - s.build_connector_meta(scheduler_output=None)

113 + # Worker reports the save copy finished → entry becomes ready.

114 + s.update_connector_output(_WorkerOutput(saves=["h1"]))

115 assert s.has_cache_item("h1") is True

116

117 # Step C: reload now works — entry is ready at update_state_after_alloc time.

118 s.update_state_after_alloc(req, 0)

119 meta_c = s.build_connector_meta(scheduler_output=None)

120 assert "h1" in meta_c.loads

106 - assert meta_c.loads["h1"] == meta_a.saves["h1"]

121 + assert _load_blocks(meta_c, "h1") == meta_a.saves["h1"]

122

123 s.shutdown()

124

125

126 def test_has_cache_item_false_when_not_consumer(monkeypatch):

112 - s = _make_scheduler(monkeypatch, is_consumer=False)

127 + s = _make_scheduler(monkeypatch, ec_role="ec_producer")

128 assert s.has_cache_item("anything") is False

129 s.shutdown()

130

@@ -123,10 +138,10 @@ def test_connector_keys_on_identifier_not_mm_hash(monkeypatch):

138 meta = s.build_connector_meta(scheduler_output=None)

139 assert "ENC_KEY" in meta.saves

140 assert "PROC_KEY" not in meta.saves

126 - # Not ready yet (1 step delay).

141 + # Not ready until the completion report.

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

143 + # After the report it becomes ready — keyed on identifier.

144 + s.update_connector_output(_WorkerOutput(saves=["ENC_KEY"]))

145 assert s.has_cache_item("ENC_KEY") is True

146 assert s.has_cache_item("PROC_KEY") is False

147 s.shutdown()

@@ -136,13 +151,11 @@ def test_cpu_region_fifo_eviction(monkeypatch):

151 # Region holds exactly 2 one-block encodings.

152 s = _make_scheduler(monkeypatch, num_blocks=2)

153

139 - # Fill with A and B, marking each ready.

154 + # Fill with A and B, marking each ready via a completion report.

155 for h in ("A", "B"):

156 s.update_state_after_alloc(_Request([_Feature(h, length=1)]), 0)

157 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

158 + s.update_connector_output(_WorkerOutput(saves=[h]))

159 assert s.has_cache_item("A") is True

160 assert s.has_cache_item("B") is True

161

@@ -161,8 +174,8 @@ def test_multiple_mm_items_per_request(monkeypatch):

174 meta = s.build_connector_meta(scheduler_output=None)

175 assert set(meta.saves) == {"h1", "h2"}

176

164 - # Step delay: mark both ready.

165 - s.build_connector_meta(scheduler_output=None)

177 + # Worker reports both saves complete.

178 + s.update_connector_output(_WorkerOutput(saves=["h1", "h2"]))

179 assert s.has_cache_item("h1") is True

180 assert s.has_cache_item("h2") is True

181

@@ -174,26 +187,6 @@ def test_multiple_mm_items_per_request(monkeypatch):

187 s.shutdown()

188

189

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 -

190 def test_loads_only_serves_hashes_touched_this_step(monkeypatch):

191 s = _make_scheduler(monkeypatch)

192 _seed_cached(s, "a", n_blocks=2)

@@ -228,10 +221,10 @@ def test_load_not_emitted_for_uncached_entry(monkeypatch):

221 s.shutdown()

222

223

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

224 +def test_load_pin_protects_blocks_until_completion(monkeypatch):

225 + """Loaded blocks stay pinned until the worker reports the load memcpy

226 + complete. Observable: a save that needs eviction cannot reclaim a

227 + still-pinned loaded entry, but can once the completion report arrives."""

228 # 2 blocks, both occupied by "a" (ready).

229 s = _make_scheduler(monkeypatch, num_blocks=2)

230 _seed_cached(s, "a", n_blocks=2)

@@ -240,94 +233,79 @@ def test_delayed_unpin_protects_blocks_during_worker_read(monkeypatch):

233 s.update_state_after_alloc(_Request([_Feature("a")]), 0)

234 meta = s.build_connector_meta(scheduler_output=None)

235 assert "a" in meta.loads

236 + (transfer_id,) = _load_ids(meta)

237

238 # Try to save "b" (needs 2 blocks, but "a" is pinned) — must fail.

239 s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)

240 meta2 = s.build_connector_meta(scheduler_output=None)

241 assert "b" not in meta2.saves # can't evict pinned "a"

242

249 - # After the delay, "a" unpins. Now "b" can evict it.

243 + # The one participant reports the load complete → "a" unpins. Now "b" can

244 + # evict it.

245 + s.update_connector_output(_WorkerOutput(loads=[transfer_id]))

246 s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)

247 meta3 = s.build_connector_meta(scheduler_output=None)

248 assert "b" in meta3.saves

249 s.shutdown()

250

251

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

252 +def _can_evict(s: ECCPUScheduler, mm_hash: str, n_blocks: int) -> bool:

253 + """Whether a new save of `mm_hash` can claim `n_blocks`, i.e. whether the

254 + blocks currently held by other entries are reclaimable."""

255 + s.update_state_after_alloc(_Request([_Feature(mm_hash, length=n_blocks)]), 0)

256 + meta = s.build_connector_meta(scheduler_output=None)

257 + return mm_hash in meta.saves

258

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

259

260 +def test_load_needs_every_participant_before_unpin(monkeypatch):

261 + """With two participating ranks, one report must not release the pin.

262

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)

263 + This is the multi-rank case: each rank copies the same CPU blocks into its

264 + own GPU memory, so the blocks stay in use until the last rank is done.

265 + """

266 + s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2)

267 _seed_cached(s, "a", n_blocks=2)

268

269 s.update_state_after_alloc(_Request([_Feature("a")]), 0)

291 - # Don't call build_connector_meta — "a" stays in _pending_loads.

270 + (transfer_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))

271 +

272 + # First rank reports: still in use by the second, so "a" must hold.

273 + s.update_connector_output(_WorkerOutput(loads=[transfer_id]))

274 + assert _can_evict(s, "b", 2) is False

275 + assert s.has_pending_push_work() is True

276 +

277 + # Second rank reports: now it is genuinely free.

278 + s.update_connector_output(_WorkerOutput(loads=[transfer_id]))

279 + assert s.has_pending_push_work() is False

280 + assert _can_evict(s, "c", 2) is True

281 s.shutdown()

282

294 - # After shutdown, scheduler reports no items (consumer role is off).

295 - assert s.has_cache_item("a") is False

283

284 +def test_late_report_does_not_release_a_later_load_of_same_hash(monkeypatch):

285 + """A straggling report from one dispatch must not release the pin taken by

286 + a subsequent dispatch of the same mm_hash.

287

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)

288 + This is what the transfer id is for: mm_hash identifies a cache entry, not

289 + a transfer, so it cannot tell the two dispatches apart.

290 + """

291 + s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2)

292 _seed_cached(s, "a", n_blocks=2)

293

304 - # Load "a" and build — moves to unpin deque.

294 + # First dispatch, fully reported by both ranks → pin released.

295 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

296 + (first_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))

297 + s.update_connector_output(_WorkerOutput(loads=[first_id, first_id]))

298 + assert s.has_pending_push_work() is False

299

300 + # Second dispatch of the same hash takes a fresh pin.

301 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)

302 + (second_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))

303 + assert second_id != first_id

304

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

305 + # A replayed report for the first dispatch must be ignored entirely.

306 + s.update_connector_output(_WorkerOutput(loads=[first_id, first_id]))

307 + assert s.has_pending_push_work() is True

308 + assert _can_evict(s, "b", 2) is False

309 s.shutdown()

310

311

@@ -350,30 +328,10 @@ def test_region_full_skips_save_and_never_blocks(monkeypatch):

328 s.shutdown()

329

330

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 -

331 def test_producer_only_never_emits_loads(monkeypatch):

332 """A producer-only scheduler must never populate meta.loads, even when

333 entries are ready."""

376 - s = _make_scheduler(monkeypatch, is_consumer=False)

334 + s = _make_scheduler(monkeypatch, ec_role="ec_producer")

335 req = _Request([_Feature("h1", length=1)])

336

337 s.update_state_after_alloc(req, 0)

@@ -381,8 +339,7 @@ def test_producer_only_never_emits_loads(monkeypatch):

339 assert "h1" in meta.saves

340 assert meta.loads == {}

341

384 - # Mark ready.

385 - s.build_connector_meta(scheduler_output=None)

342 + s.update_connector_output(_WorkerOutput(saves=["h1"]))

343

344 # Even if we re-encounter the same feature, no load.

345 s.update_state_after_alloc(req, 0)

@@ -393,7 +350,7 @@ def test_producer_only_never_emits_loads(monkeypatch):

350

351 def test_consumer_only_never_emits_saves(monkeypatch):

352 """A consumer-only scheduler must never populate meta.saves."""

396 - s = _make_scheduler(monkeypatch, is_producer=False)

353 + s = _make_scheduler(monkeypatch, ec_role="ec_consumer")

354 _seed_cached(s, "a", n_blocks=1)

355

356 req = _Request([_Feature("a", length=1)])

@@ -406,7 +363,7 @@ def test_consumer_only_never_emits_saves(monkeypatch):

363

364 def test_consumer_only_has_cache_item(monkeypatch):

365 """Consumer-only: has_cache_item works normally for ready entries."""

409 - s = _make_scheduler(monkeypatch, is_producer=False)

366 + s = _make_scheduler(monkeypatch, ec_role="ec_consumer")

367 assert s.has_cache_item("a") is False

368 _seed_cached(s, "a", n_blocks=1)

369 assert s.has_cache_item("a") is True

@@ -431,50 +388,75 @@ def test_save_not_emitted_for_already_cached_entry(monkeypatch):

388 s.shutdown()

389

390

434 -# ── first-finish fast-path (idle engine) ────────────────────────────────────

391 +# ── completion reporting (update_connector_output) ──────────────────────────

392

393

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

394 +def test_update_connector_output_ignores_foreign_meta(monkeypatch):

395 + """A worker output whose payload is not an ECCPUWorkerMetadata is a

396 + no-op — no cache mutation, no crash."""

397 + s = _make_scheduler(monkeypatch)

398 + _seed_cached(s, "a", n_blocks=1)

399

443 - s.update_state_after_alloc(req, 0)

400 + class _Foreign:

401 + ec_connector_worker_meta = object()

402 +

403 + s.update_connector_output(_Foreign()) # must not raise

404 + assert s.has_cache_item("a") is True

405 + s.shutdown()

406 +

407 +

408 +def test_update_connector_output_ignores_unknown_report(monkeypatch):

409 + """Reports for a never-seen save hash or a transfer id this scheduler never

410 + dispatched are dropped rather than mutating unrelated state."""

411 + s = _make_scheduler(monkeypatch)

412 + s.update_connector_output(_WorkerOutput(saves=["gone"], loads=[4242]))

413 + assert s.has_cache_item("gone") is False

414 + s.shutdown()

415 +

416 +

417 +# ── has_pending_push_work ───────────────────────────────────────────────────

418 +

419 +

420 +def test_has_pending_push_work_tracks_inflight_save(monkeypatch):

421 + """A dispatched-but-unconfirmed save keeps push work pending until the

422 + completion report marks it ready."""

423 + s = _make_scheduler(monkeypatch)

424 + assert s.has_pending_push_work() is False

425 +

426 + s.update_state_after_alloc(_Request([_Feature("h1", length=1)]), 0)

427 s.build_connector_meta(scheduler_output=None)

445 - assert s.has_cache_item("h1") is False

428 + assert s.has_pending_push_work() is True # save not yet confirmed

429

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

430 + s.update_connector_output(_WorkerOutput(saves=["h1"]))

431 + assert s.has_pending_push_work() is False

432 s.shutdown()

433

434

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)

435 +def test_has_pending_push_work_tracks_inflight_load(monkeypatch):

436 + """A dispatched-but-unconfirmed load keeps push work pending until the

437 + unpin completion report arrives."""

438 + s = _make_scheduler(monkeypatch)

439 + _seed_cached(s, "a", n_blocks=1)

440 + assert s.has_pending_push_work() is False # ready + unpinned

441

459 - req = _Request([_Feature("a")], request_id="req_y")

460 - s.update_state_after_alloc(req, 0)

442 + s.update_state_after_alloc(_Request([_Feature("a")]), 0)

443 meta = s.build_connector_meta(scheduler_output=None)

462 - assert "a" in meta.loads

444 + assert s.has_pending_push_work() is True # pinned, awaiting unpin

445

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

446 + s.update_connector_output(_WorkerOutput(loads=_load_ids(meta)))

447 + assert s.has_pending_push_work() is False

448 + s.shutdown()

449

471 - # req_y finishes — unpin fires in build_connector_meta (fast-path).

472 - s.build_connector_meta(_SchedulerOutput(finished_req_ids={"req_y"}))

450

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

451 +# ── shutdown ────────────────────────────────────────────────────────────────

452 +

453 +

454 +def test_shutdown_disables_roles_and_cleans_region(monkeypatch):

455 + """After shutdown the scheduler serves no items and the region is

456 + cleaned up."""

457 + s = _make_scheduler(monkeypatch, num_blocks=2)

458 + _seed_cached(s, "a", n_blocks=2)

459 + assert s.has_cache_item("a") is True

460 +

461 s.shutdown()

462 + assert s.has_cache_item("a") is False # consumer role disabled

tests/v1/ec_connector/unit/cpu/scheduler/test_step_tracker.py deleted

-162

@@ -1,162 +0,0 @@

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/worker/test_worker.py

+342 -71

@@ -22,14 +22,16 @@ Mocking policy

23 import contextlib

24 import logging

25 +import time

26 import uuid

27 +from collections import deque

28 from unittest.mock import MagicMock, Mock, patch

30 +import numpy as np

31 import pytest

32 import torch

31 -from vllm.config import VllmConfig

32 -from vllm.config.parallel import ParallelConfig

34 +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config

35 from vllm.distributed.ec_transfer.ec_connector.cpu.common import (

36 ECCPUConnectorMetadata,

37 )

@@ -55,6 +57,18 @@ _requires_accelerator = pytest.mark.skipif(

57 reason="exercises real accelerator stream/event coordination in ECCPUWorker",

58 )

60 +# The memory-lifetime tests hold a copy back with `torch.cuda._sleep`, which the

61 +# CUDA-like platforms provide but XPU does not.

62 +_requires_cuda_alike = pytest.mark.skipif(

63 + not current_platform.is_cuda_alike(),

64 + reason="stalling the compute stream requires a CUDA-like platform",

65 +)

66 +

67 +# Cycles to stall the compute stream for. Must comfortably outlast the host-side

68 +# work the lifetime tests do while a copy is deliberately held back — ~0.5 s on

69 +# a current datacenter GPU.

70 +_STALL_CYCLES = 1_000_000_000

71 +

73 def _make_region() -> ECSharedRegion:

74 """Fresh region backed by a real per-test mmap file."""

@@ -66,18 +80,28 @@ def _make_region() -> ECSharedRegion:

82 def _vllm_config(rank: int = 0) -> Mock:

69 - cfg = Mock(spec=VllmConfig)

70 - cfg.parallel_config = Mock(spec=ParallelConfig)

71 - cfg.parallel_config.rank = rank

72 - cfg.model_config = Mock()

73 - cfg.model_config.dtype = _DTYPE

74 - return cfg

83 + return create_ec_vllm_config(rank=rank, dtype=_DTYPE)

86 def _meta(

87 *, saves: dict | None = None, loads: dict | None = None

88 ) -> ECCPUConnectorMetadata:

80 - return ECCPUConnectorMetadata(saves=saves or {}, loads=loads or {})

89 + """Build step metadata from plain `{mm_hash: block_ids}` dicts.

90 +

91 + Load transfer ids are synthesized here so tests can keep naming loads by

92 + mm_hash; `_load_id` recovers the id the worker will report.

93 + """

94 + return ECCPUConnectorMetadata(

95 + saves=saves or {},

96 + loads={

97 + mm_hash: (idx, block_ids)

98 + for idx, (mm_hash, block_ids) in enumerate((loads or {}).items())

99 + },

100 + )

101 +

102 +

103 +def _load_id(meta: ECCPUConnectorMetadata, mm_hash: str) -> int:

104 + return meta.loads[mm_hash][0]

105

106

107 # ── fixtures ─────────────────────────────────────────────────────────────────

@@ -136,19 +160,52 @@ def make_worker():

160 region.cleanup()

161

162

163 +def _wait_for_completion(

164 + worker: ECCPUWorker, expected, direction: str, timeout_s: float = 30.0

165 +) -> None:

166 + """Poll ``build_connector_worker_meta`` until ``expected`` is reported done.

167 +

168 + Saves are reported by mm_hash, loads by transfer id.

169 + """

170 + deadline = time.monotonic() + timeout_s

171 + while time.monotonic() < deadline:

172 + meta = worker.build_connector_worker_meta()

173 + if meta is not None and expected in getattr(meta, f"completed_{direction}"):

174 + return

175 + time.sleep(0.001)

176 + raise AssertionError(f"{direction} for {expected!r} never reported complete")

177 +

178 +

179 +def _warm_up_stream_pools(worker: ECCPUWorker) -> None:

180 + """Run one throwaway save so the stream and event pools are populated.

181 +

182 + The first CUDA stream creation in a process blocks the host for ~0.5 s on

183 + driver initialization, which would swallow the stall the lifetime tests

184 + below install. One flush puts a stream and its events in the worker's pools,

185 + matching the steady state every step after the first sees.

186 + """

187 + scratch = torch.zeros(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)

188 + worker.save_caches({"w": scratch}, "w", _meta(saves={"w": [0]}))

189 + worker.flush_saves()

190 + _wait_for_completion(worker, "w", "saves")

191 + assert worker._stream_pool, "warm-up did not recycle a stream"

192 +

193 +

194 # ── save_caches ──────────────────────────────────────────────────────────────

195

196

197 @_requires_accelerator

198 @pytest.mark.parametrize(

144 - "n_elements,n_blocks",

199 + "n_elements,block_ids",

200 [

146 - (3 * _HIDDEN_DIM, 3), # exact fit: 3 full blocks, no remainder

147 - (3 * _HIDDEN_DIM + 4, 4), # partial last block: 4 fp16 trail in slot 3

201 + (3 * _HIDDEN_DIM, [7, 2, 5]), # exact fit: 3 full blocks, no remainder

202 + (3 * _HIDDEN_DIM + 4, [7, 2, 5, 0]), # partial last block: 4 fp16 trail

203 + (3 * _HIDDEN_DIM, [2, 3, 4]), # adjacent: one coalesced copy

204 + (3 * _HIDDEN_DIM + 4, [2, 3, 4, 6]), # coalesced run, then a gap

205 ],

149 - ids=["exact-fit", "partial-last-block"],

206 + ids=["exact-fit", "partial-last-block", "adjacent-blocks", "run-then-gap"],

207 )

151 -def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks):

208 +def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, block_ids):

209 """``save_caches`` + ``flush_saves`` copies the source GPU tensor's bytes

210 into the block IDs named by ``meta.saves``, in the order they appear,

211 via a single batched GPU→CPU copy.

@@ -156,6 +213,8 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks

213 When ``total_bytes`` is not a multiple of ``block_size_bytes`` the last

214 block is partially written; the unwritten tail must remain whatever

215 was there before. Blocks not named in ``meta.saves`` must be untouched.

216 + Blocks adjacent in the region share one descriptor, which must land the

217 + same bytes in the same blocks as one descriptor per block does.

218 """

219 worker = make_worker()

220 sentinel = 0x5A

@@ -165,9 +224,9 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks

224 expected_bytes = src.cpu().reshape(-1).view(torch.uint8)

225 total_bytes = n_elements * _DTYPE.itemsize

226

168 - block_ids = [7, 2, 5, 0][:n_blocks]

227 worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))

228 worker.flush_saves()

229 + _wait_for_completion(worker, "h", "saves")

230

231 for slot, block_idx in enumerate(block_ids):

232 block_byte_start = slot * _BLOCK_SIZE_BYTES

@@ -189,6 +248,42 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks

248 )

249

250

251 +@_requires_accelerator

252 +@pytest.mark.parametrize(

253 + "n_elements,block_ids,expected_sizes",

254 + [

255 + (3 * _HIDDEN_DIM, [2, 3, 4], [3 * _BLOCK_SIZE_BYTES]),

256 + (3 * _HIDDEN_DIM, [7, 2, 5], [_BLOCK_SIZE_BYTES] * 3),

257 + (3 * _HIDDEN_DIM, [2, 3, 6], [2 * _BLOCK_SIZE_BYTES, _BLOCK_SIZE_BYTES]),

258 + (2 * _HIDDEN_DIM + 4, [2, 3, 4], [(2 * _HIDDEN_DIM + 4) * _DTYPE.itemsize]),

259 + (_HIDDEN_DIM, [2, 3, 4], [_BLOCK_SIZE_BYTES]),

260 + ],

261 + ids=[

262 + "one-run",

263 + "no-run",

264 + "run-then-gap",

265 + "run-stops-at-output",

266 + "blocks-past-output-skipped",

267 + ],

268 +)

269 +def test_save_caches_coalesces_only_adjacent_blocks(

270 + make_worker, n_elements, block_ids, expected_sizes

271 +):

272 + """One descriptor per run of region-adjacent blocks, not one per block.

273 +

274 + Per-descriptor overhead dominates these copies, so the byte tests above

275 + would pass either way; this pins the descriptor layout that makes them

276 + fast. A coalesced copy also stops at the encoder output, since entries are

277 + sized by placeholder count, which can exceed the number of embeddings.

278 + """

279 + worker = make_worker()

280 + src = torch.arange(n_elements, dtype=_DTYPE, device=DEVICE_TYPE)

281 +

282 + worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))

283 +

284 + assert worker._save_bufs.sizes[: worker._save_count].tolist() == expected_sizes

285 +

286 +

287 def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker):

288 """When the scheduler hasn't pre-allocated blocks for ``mm_hash``,

289 ``save_caches`` + ``flush_saves`` is a pure no-op."""

@@ -245,6 +340,7 @@ def test_save_caches_batches_multiple_hashes(make_worker):

340 worker.save_caches(cache, "a", _meta(saves={"a": [1], "b": [3]}))

341 worker.save_caches(cache, "b", _meta(saves={"a": [1], "b": [3]}))

342 worker.flush_saves()

343 + _wait_for_completion(worker, "a", "saves")

344

345 expected_a = src_a.cpu().view(torch.uint8)

346 expected_b = src_b.cpu().view(torch.uint8)

@@ -258,6 +354,50 @@ def test_save_caches_batches_multiple_hashes(make_worker):

354 assert torch.all(worker._region.blocks[idx] == sentinel)

355

356

357 +class _HighAddress:

358 + """A device tensor reporting an address with the high bit set.

359 +

360 + XPU USM allocations live up there and no host allocator will hand out such

361 + an address, so naming one is the only way to reach the arithmetic. Every

362 + attribute other than the address comes from the real tensor.

363 + """

364 +

365 + def __init__(self, real: torch.Tensor, address: int):

366 + self._real = real

367 + self._address = address

368 +

369 + def __getattr__(self, name):

370 + return getattr(self._real, name)

371 +

372 + def view(self, *args):

373 + return self

374 +

375 + def data_ptr(self) -> int:

376 + return self._address

377 +

378 +

379 +@_requires_accelerator

380 +def test_save_descriptors_hold_high_bit_source_addresses(make_worker):

381 + """Source addresses >= 2**63 survive into the descriptor buffers."""

382 + worker = make_worker()

383 + real = torch.zeros(2 * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)

384 + address = 2**64 - 8 * _BLOCK_SIZE_BYTES

385 +

386 + # Two runs, so the second descriptor's address is an offset off the base.

387 + worker.save_caches(

388 + {"a": _HighAddress(real, address)}, "a", _meta(saves={"a": [1, 5]})

389 + )

390 +

391 + assert worker._save_count == 2

392 + stored = worker._save_bufs.src_np[:2].astype(np.uint64)

393 + assert list(stored) == [address, address + _BLOCK_SIZE_BYTES]

394 +

395 + # The batch describes an address nothing may read from; drop it unflushed.

396 + worker._save_bufs = None

397 + worker._save_count = 0

398 + worker._save_mm_hashes = []

399 +

400 +

401 # ── start_load_caches ────────────────────────────────────────────────────────

402

403

@@ -288,19 +428,31 @@ def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker

428

429

430 @_requires_accelerator

291 -def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker):

292 - """If ``encoder_cache`` already holds the ``mm_hash``, the worker must

293 - not overwrite it."""

431 +@pytest.mark.parametrize(

432 + "loads,expected_sizes",

433 + [

434 + ({"a": [2, 3], "b": [4, 5]}, [4 * _BLOCK_SIZE_BYTES]),

435 + ({"b": [4, 5], "a": [2, 3]}, [4 * _BLOCK_SIZE_BYTES]),

436 + ({"a": [2, 3], "b": [6, 7]}, [2 * _BLOCK_SIZE_BYTES] * 2),

437 + ],

438 + ids=["adjacent-entries", "adjacent-entries-reversed", "gap-between-entries"],

439 +)

440 +def test_start_load_caches_coalesces_across_entries(make_worker, loads, expected_sizes):

441 + """Blocks adjacent in the region share a descriptor across entries too.

442 +

443 + Entries are ordered by first block, so whether they coalesce follows from

444 + the region layout rather than the order the scheduler dispatched them.

445 + """

446 worker = make_worker()

295 - worker._region.blocks[0].fill_(0x42)

447 + captured: list[list[int]] = []

448

297 - sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device=DEVICE_TYPE)

298 - encoder_cache = {"h": sentinel}

299 - worker.start_load_caches(encoder_cache, _meta(loads={"h": [0]}))

449 + with patch(

450 + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.swap_blocks_batch",

451 + side_effect=lambda src, dst, sizes, **kw: captured.append(sizes.tolist()),

452 + ):

453 + worker.start_load_caches({}, _meta(loads=loads))

454

301 - assert encoder_cache["h"] is sentinel, (

302 - "existing encoder_cache entry must not be replaced"

303 - )

455 + assert captured == [expected_sizes]

456

457

458 @_requires_accelerator

@@ -315,30 +467,34 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker):

467

468 @_requires_accelerator

469 def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):

318 - """Cached entries are preserved while new ones are loaded."""

470 + """Every hash in ``meta.loads`` is copied from mmap, including one whose

471 + key is already resident in ``encoder_cache`` — there is no worker-side

472 + skip."""

473 worker = make_worker()

320 - n_blocks = 2

474 + n_blocks = 3

475 src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape(

476 n_blocks, _HIDDEN_DIM

477 )

478 src_int8 = src_orig.view(torch.int8).reshape(n_blocks, _BLOCK_SIZE_BYTES)

325 - new_block_ids = [4, 2]

326 - for i, idx in enumerate(new_block_ids):

479 + block_ids = [0, 4, 2]

480 + for i, idx in enumerate(block_ids):

481 worker._region.blocks[idx].copy_(src_int8[i])

482

483 cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device=DEVICE_TYPE)

484 encoder_cache: dict[str, torch.Tensor] = {"cached_h": cached_tensor}

485 worker.start_load_caches(

486 encoder_cache,

333 - _meta(loads={"cached_h": [0], "new_h": new_block_ids}),

487 + _meta(loads={"resident_h": [block_ids[0]], "new_h": block_ids[1:]}),

488 )

489

336 - assert encoder_cache["cached_h"] is cached_tensor

337 - assert "new_h" in encoder_cache

338 - out = encoder_cache["new_h"]

339 - assert out.shape == (n_blocks, _HIDDEN_DIM)

340 - assert out.dtype == _DTYPE

341 - assert torch.equal(out.cpu(), src_orig)

490 + resident = encoder_cache["resident_h"]

491 + assert resident is not cached_tensor, "resident entry must be reloaded, not skipped"

492 + assert torch.equal(resident.cpu(), src_orig[:1])

493 +

494 + new = encoder_cache["new_h"]

495 + assert new.shape == (2, _HIDDEN_DIM)

496 + assert new.dtype == _DTYPE

497 + assert torch.equal(new.cpu(), src_orig[1:])

498

499

500 @_requires_accelerator

@@ -382,6 +538,7 @@ def test_save_then_load_round_trips_bytes(make_worker):

538 ).reshape(n_blocks, _HIDDEN_DIM)

539 worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))

540 worker.flush_saves()

541 + _wait_for_completion(worker, "h", "saves")

542

543 encoder_cache: dict[str, torch.Tensor] = {}

544 worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids}))

@@ -392,18 +549,133 @@ def test_save_then_load_round_trips_bytes(make_worker):

549 assert torch.equal(out.cpu(), src.cpu())

550

551

552 +# ── memory lifetime across in-flight copies ─────────────────────────────────

553 +

554 +

555 +@_requires_cuda_alike

556 +def test_save_survives_encoder_cache_free_before_copy_runs(make_worker):

557 + """The bytes handed to ``save_caches`` must reach the mmap intact even when

558 + the caller drops its ``encoder_cache`` reference before the copy has run.

559 +

560 + The model runner pops encoder cache entries on whatever step the scheduler

561 + asks it to, with no dependency on the save copy having landed. Dropping the

562 + last reference hands the GPU memory back to the caching allocator, which is

563 + free to reuse it for the next same-sized allocation on the compute stream —

564 + the copy reads that memory on a different stream, so the allocator has no

565 + reason to hold it back.

566 +

567 + The stall makes the hazard deterministic rather than timing-dependent: the

568 + save copy is gated behind the compute stream, so it provably has not started

569 + while the overwrite (issued on an ungated stream and waited on) completes.

570 +

571 + The mmap must be pinned, as it is in production: a copy into pageable host

572 + memory blocks the host until it lands, which would hide the hazard.

573 + """

574 + worker = make_worker(pin_memory_available=True)

575 + _warm_up_stream_pools(worker)

576 + worker._region.blocks.fill_(0x5A)

577 +

578 + block_ids = [0, 1, 2]

579 + saved = torch.full(

580 + (len(block_ids), _HIDDEN_DIM), 1.0, dtype=_DTYPE, device=DEVICE_TYPE

581 + )

582 + expected = saved.cpu().reshape(-1).view(torch.uint8)

583 +

584 + encoder_cache = {"h": saved}

585 + worker.save_caches(encoder_cache, "h", _meta(saves={"h": block_ids}))

586 +

587 + torch.cuda._sleep(_STALL_CYCLES)

588 + worker.flush_saves()

589 +

590 + saved_ptr = saved.data_ptr()

591 + del saved

592 + encoder_cache.clear()

593 +

594 + # Same size on the same stream, so the allocator hands back the block just

595 + # freed. Overwriting it from an ungated stream — and waiting for that write

596 + # — orders the overwrite strictly before the still-stalled save copy.

597 + overwrite = torch.empty(

598 + (len(block_ids), _HIDDEN_DIM), dtype=_DTYPE, device=DEVICE_TYPE

599 + )

600 + reused = overwrite.data_ptr() == saved_ptr

601 + other_stream = torch.cuda.Stream()

602 + with torch.cuda.stream(other_stream):

603 + overwrite.fill_(-1.0)

604 + other_stream.synchronize()

605 +

606 + _wait_for_completion(worker, "h", "saves")

607 +

608 + actual = worker._region.blocks[block_ids].reshape(-1).view(torch.uint8)

609 + assert torch.equal(actual, expected), (

610 + "mmap holds bytes written after save_caches was called; the copy read "

611 + f"freed source memory (allocator reused the freed block: {reused})"

612 + )

613 +

614 +

615 +@_requires_cuda_alike

616 +def test_load_buffer_survives_eviction_while_consumer_read_is_queued(make_worker):

617 + """A loaded encoder cache entry must keep its bytes for a consumer already

618 + queued on the compute stream, even after the entry is dropped and another

619 + load is dispatched.

620 +

621 + ``start_load_caches`` allocates its destination inside the load stream's

622 + context, which ties that memory to the load stream in the caching allocator.

623 + Dropping the encoder cache entry returns it to the load stream's pool, where

624 + the next load can claim it — while the model's read of it is still queued on

625 + the compute stream.

626 + """

627 + worker = make_worker(pin_memory_available=True)

628 + _warm_up_stream_pools(worker)

629 + a_blocks, b_blocks = [4, 5], [6, 7]

630 + for block_idx, fill in zip(a_blocks + b_blocks, (0x11, 0x12, 0x21, 0x22)):

631 + worker._region.blocks[block_idx].fill_(fill)

632 + expected_a = worker._region.blocks[a_blocks].reshape(-1).clone()

633 +

634 + encoder_cache: dict[str, torch.Tensor] = {}

635 + meta_a = _meta(loads={"a": a_blocks})

636 + worker.start_load_caches(encoder_cache, meta_a)

637 + loaded_a = encoder_cache["a"]

638 +

639 + # Queue a consumer of the loaded entry on the compute stream behind a stall,

640 + # the way the model reads the encoder cache later in the step.

641 + torch.cuda._sleep(_STALL_CYCLES)

642 + consumed = loaded_a.clone()

643 +

644 + # Draining the completion report recycles the load stream into the pool, so

645 + # the next load allocates from the same allocator pool the entry was freed to.

646 + _wait_for_completion(worker, _load_id(meta_a, "a"), "loads")

647 +

648 + loaded_a_ptr = loaded_a.data_ptr()

649 + del loaded_a

650 + encoder_cache.clear()

651 +

652 + worker.start_load_caches(encoder_cache, _meta(loads={"b": b_blocks}))

653 + reused = encoder_cache["b"].data_ptr() == loaded_a_ptr

654 + torch.accelerator.synchronize()

655 +

656 + assert torch.equal(consumed.cpu().reshape(-1).view(torch.int8), expected_a), (

657 + "queued consumer read bytes from a later load; the load destination was "

658 + f"recycled while still in use (allocator reused the buffer: {reused})"

659 + )

660 +

661 +

662 # ── buffer recycling ────────────────────────────────────────────────────────

663

664

665 @_requires_accelerator

666 def test_buffer_pool_is_reused_across_save_steps(make_worker):

400 - """After flush_saves, descriptor buffers are returned to the pool and

401 - reused on the next flush — no reallocation."""

667 + """Once a save copy completes its descriptor buffers return to the pool and

668 + are reused by the next flush — no reallocation.

669 +

670 + The buffers hold the addresses the copy reads, so they stay out of the pool

671 + until the transfer's end event fires.

672 + """

673 worker = make_worker()

674 src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)

675

676 worker.save_caches({"h": src}, "h", _meta(saves={"h": [0]}))

677 worker.flush_saves()

678 + _wait_for_completion(worker, "h", "saves")

679

680 assert len(worker._buf_pool._pool) == 1

681 buf_id = id(worker._buf_pool._pool[0].src_ptrs)

@@ -411,6 +683,7 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker):

683 # Second step reuses the same buffer.

684 worker.save_caches({"h": src}, "h", _meta(saves={"h": [1]}))

685 worker.flush_saves()

686 + _wait_for_completion(worker, "h", "saves")

687

688 assert len(worker._buf_pool._pool) == 1

689 assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id

@@ -418,35 +691,29 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker):

691

692 @_requires_accelerator

693 def test_buffer_pool_is_reused_across_load_steps(make_worker):

421 - """After start_load_caches, descriptor buffers are returned to the pool

422 - and reused on the next call."""

694 + """Once a load copy completes its descriptor buffers return to the pool and

695 + are reused by the next call — no reallocation."""

696 worker = make_worker()

697 worker._region.blocks[0].fill_(0x01)

698 worker._region.blocks[1].fill_(0x02)

699

700 encoder_cache: dict[str, torch.Tensor] = {}

428 - worker.start_load_caches(encoder_cache, _meta(loads={"a": [0]}))

701 + meta = _meta(loads={"a": [0]})

702 + worker.start_load_caches(encoder_cache, meta)

703 + _wait_for_completion(worker, _load_id(meta, "a"), "loads")

704

705 assert len(worker._buf_pool._pool) == 1

706 buf_id = id(worker._buf_pool._pool[0].src_ptrs)

707

708 encoder_cache2: dict[str, torch.Tensor] = {}

434 - worker.start_load_caches(encoder_cache2, _meta(loads={"b": [1]}))

709 + meta2 = _meta(loads={"b": [1]})

710 + worker.start_load_caches(encoder_cache2, meta2)

711 + _wait_for_completion(worker, _load_id(meta2, "b"), "loads")

712

713 assert len(worker._buf_pool._pool) == 1

714 assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id

715

716

440 -# ── stream management ────────────────────────────────────────────────────────

441 -

442 -

443 -@_requires_accelerator

444 -def test_stream_initialized_at_construction(make_worker):

445 - """``_load_stream`` must be a fully initialized accelerator stream."""

446 - worker = make_worker()

447 - assert isinstance(worker._load_stream, current_platform.Stream)

448 -

449 -

717 # ── lifecycle ────────────────────────────────────────────────────────────────

718

719

@@ -500,12 +767,15 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):

767 worker = object.__new__(ECCPUWorker)

768 mock_region = Mock(spec=ECSharedRegion)

769 worker._region = mock_region

503 - worker._load_stream = MagicMock()

770 worker._save_bufs = None

771 worker._save_count = 0

772 + worker._save_mm_hashes = []

773 + worker._inflight_saves = deque()

774 + worker._inflight_loads = deque()

775 + worker._stream_pool = []

776 + worker._event_pool = []

777

778 worker.shutdown()

508 - worker._load_stream.synchronize.assert_called_once()

779 mock_region.cleanup.assert_called_once()

780

781 mock_region.cleanup.side_effect = RuntimeError("boom")

@@ -525,12 +795,13 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):

795

796 @_requires_accelerator

797 def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):

528 - """Full pipeline: scheduler allocates blocks, worker saves GPU tensor to

529 - mmap via flush_saves, scheduler marks ready after step delay, worker

530 - loads from mmap back to GPU, and the result matches the original.

798 + """Full pipeline: scheduler allocates blocks, worker saves a GPU tensor to

799 + mmap via flush_saves, the worker's completion report marks the entry ready,

800 + the worker loads from mmap back to GPU, and the result matches the original.

801

802 Exercises the real scheduler + worker cooperation through a shared

533 - ECSharedRegion, with real accelerator transfers and stream coordination.

803 + ECSharedRegion, with real accelerator transfers and stream coordination,

804 + and the event-driven mark_ready path.

805 """

806 import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod

807 from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import (

@@ -544,15 +815,7 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):

815 # Build scheduler sharing the same region.

816 monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)

817

547 - class _EC:

548 - is_ec_producer = True

549 - is_ec_consumer = True

550 -

551 - class _Cfg:

552 - ec_transfer_config = _EC()

553 - max_concurrent_batches = 1

554 -

555 - scheduler = ECCPUScheduler(_Cfg())

818 + scheduler = ECCPUScheduler(_vllm_config())

819

820 # -- Step 1: scheduler allocates, worker saves --

821 n_blocks = 3

@@ -580,16 +843,24 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):

843 worker.save_caches(encoder_cache, "img_001", meta_save)

844 worker.flush_saves()

845

583 - # -- Step 2: readiness delay (max_concurrent_batches=1) --

584 - meta_step2 = scheduler.build_connector_meta(scheduler_output=None)

846 + # -- Step 2: worker reports the save memcpy complete → scheduler marks ready --

847 + torch.accelerator.synchronize()

848 + worker_meta = worker.build_connector_worker_meta()

849 + assert worker_meta is not None

850 + assert "img_001" in worker_meta.completed_saves

851 +

852 + class _Output:

853 + ec_connector_worker_meta = worker_meta

854 +

855 + scheduler.update_connector_output(_Output())

856 assert scheduler.has_cache_item("img_001") is True

586 - assert meta_step2.loads == {} # no load requested yet

857

858 # -- Step 3: scheduler emits load, worker loads --

859 scheduler.update_state_after_alloc(_Request(), 0)

860 meta_load = scheduler.build_connector_meta(scheduler_output=None)

861 assert "img_001" in meta_load.loads

592 - assert meta_load.loads["img_001"] == meta_save.saves["img_001"]

862 + _, load_blocks = meta_load.loads["img_001"]

863 + assert load_blocks == meta_save.saves["img_001"]

864

865 load_cache: dict[str, torch.Tensor] = {}

866 worker.start_load_caches(load_cache, meta_load)

tests/v1/ec_connector/unit/test_ec_cpu_connector.py

+4 -3

@@ -48,9 +48,10 @@ def _make_unique_image(base_image, iteration: int):

48 def _wait_for_ec_ready(llm: LLM) -> None:

49 """Send a dummy request to force a scheduler step.

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.

51 + When generate() returns, the save copy has been issued to the GPU but its

52 + completion event may fire on a later step, so the entry is not yet marked

53 + ready. This extra step lets the scheduler poll the worker's completion

54 + report and mark the entry ready.

55 """

56 llm.generate("hi", SamplingParams(max_tokens=1), use_tqdm=False)

tests/v1/ec_connector/unit/test_metadata.py

+13 -13

@@ -38,19 +38,19 @@ def test_metadata_saves_dict_operations():

38 assert meta.saves["mm_hash_2"] == [10, 11]

41 -def test_metadata_loads_dict_operations():

42 - """The loads dict supports standard dict operations."""

41 +def test_metadata_loads_carry_a_transfer_id_and_blocks():

42 + """Each load entry is a (transfer_id, block_ids) pair, so the worker can

43 + report the transfer rather than the mm_hash."""

44 meta = ECCPUConnectorMetadata()

45 - # Add multiple entries

46 - meta.loads["mm_hash_1"] = [5, 6, 7]

47 - meta.loads["mm_hash_2"] = [100]

46 + meta.loads["mm_hash_1"] = (0, [5, 6, 7])

47 + meta.loads["mm_hash_2"] = (1, [100])

49 - assert "mm_hash_1" in meta.loads

50 - assert "mm_hash_2" in meta.loads

49 assert len(meta.loads) == 2

52 - assert meta.loads["mm_hash_1"] == [5, 6, 7]

53 - assert meta.loads["mm_hash_2"] == [100]

50 + assert meta.loads["mm_hash_1"] == (0, [5, 6, 7])

51 + transfer_id, block_ids = meta.loads["mm_hash_2"]

52 + assert transfer_id == 1

53 + assert block_ids == [100]

56 def test_metadata_saves_and_loads_are_independent():

@@ -58,11 +58,11 @@ def test_metadata_saves_and_loads_are_independent():

58 meta = ECCPUConnectorMetadata()

60 meta.saves["key"] = [1, 2]

61 - meta.loads["key"] = [3, 4]

61 + meta.loads["key"] = (0, [3, 4])

63 # Same key but different dicts should not interfere

64 assert meta.saves["key"] == [1, 2]

65 - assert meta.loads["key"] == [3, 4]

65 + assert meta.loads["key"] == (0, [3, 4])

67 # Modifying one should not affect the other

68 meta.saves["other"] = [99]

@@ -74,11 +74,11 @@ def test_metadata_block_indices_are_integers():

74 meta = ECCPUConnectorMetadata()

76 meta.saves["hash"] = [0, 100, 999, 65535]

77 - meta.loads["hash"] = [1, 2, 3]

77 + meta.loads["hash"] = (0, [1, 2, 3])

79 for idx in meta.saves["hash"]:

80 assert isinstance(idx, int)

81 - for idx in meta.loads["hash"]:

81 + for idx in meta.loads["hash"][1]:

82 assert isinstance(idx, int)

tests/v1/ec_connector/unit/utils.py new

+70

@@ -0,0 +1,70 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Shared config builders for EC connector unit tests."""

4 +

5 +import uuid

6 +from unittest.mock import Mock

7 +

8 +import torch

9 +

10 +from vllm.config import ModelConfig, VllmConfig

11 +from vllm.config.ec_transfer import ECRole, ECTransferConfig

12 +from vllm.config.parallel import ParallelConfig

13 +

14 +# The scheduler derives its per-load ack count from the PCP size only under the

15 +# `mp` backend -- which also keeps the world size off the local GPU count, so

16 +# TP > 1 stays expressible on a single-GPU host.

17 +_EXECUTOR_BACKEND = "mp"

18 +_PCP_SIZE = 1

19 +

20 +

21 +def create_ec_vllm_config(

22 + *,

23 + ec_role: ECRole = "ec_both",

24 + tensor_parallel_size: int = 1,

25 + rank: int = 0,

26 + dtype: torch.dtype = torch.float16,

27 +) -> Mock:

28 + """Build a `VllmConfig` stand-in for EC connector unit tests.

29 +

30 + `ec_transfer_config` and `parallel_config` are the real config objects, so

31 + role derivation and the rank arithmetic the connector depends on behave as

32 + they do in production.

33 +

34 + `model_config` is a stub because nothing under test needs a real one: the

35 + connector reads only `dtype` here, and the paths that inspect the model

36 + itself are not exercised. Building a real `ModelConfig` would resolve an HF

37 + model and inspect its architecture.

38 +

39 + Args:

40 + ec_role: EC role this instance plays. `is_ec_producer` and

41 + `is_ec_consumer` are derived from it, as in production.

42 + tensor_parallel_size: TP degree.

43 + rank: Global rank of this worker.

44 + dtype: Encoder cache dtype.

45 +

46 + Returns:

47 + A `VllmConfig`-specced mock carrying the configs above.

48 + """

49 + parallel_config = ParallelConfig(

50 + tensor_parallel_size=tensor_parallel_size,

51 + prefill_context_parallel_size=_PCP_SIZE,

52 + distributed_executor_backend=_EXECUTOR_BACKEND,

53 + rank=rank,

54 + )

55 + ec_transfer_config = ECTransferConfig(

56 + # The role properties report False unless a connector is named.

57 + ec_connector="ECCPUConnector",

58 + ec_role=ec_role,

59 + engine_id=str(uuid.uuid4()),

60 + )

61 +

62 + model_config = Mock(spec=ModelConfig)

63 + model_config.dtype = dtype

64 +

65 + vllm_config = Mock(spec=VllmConfig)

66 + vllm_config.ec_transfer_config = ec_transfer_config

67 + vllm_config.parallel_config = parallel_config

68 + vllm_config.model_config = model_config

69 + vllm_config.instance_id = f"ec-test-{uuid.uuid4()}"

70 + return vllm_config

vllm/distributed/ec_transfer/ec_connector/cpu/common.py

+36 -5

@@ -7,7 +7,10 @@ from typing import TYPE_CHECKING

8 import torch

10 -from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata

10 +from vllm.distributed.ec_transfer.ec_connector.base import (

11 + ECConnectorMetadata,

12 + ECConnectorWorkerMetadata,

13 +)

14 from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (

15 ECSharedRegion,

16 )

@@ -29,10 +32,38 @@ class ECCPUConnectorMetadata(ECConnectorMetadata):

32 # encoder_cache[mm_hash] → mmap at these block IDs.

33 saves: dict[str, list[int]] = field(default_factory=dict)

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)

35 + # Consumer role: mm_hashes whose bytes are available in the local mmap,

36 + # mapped to (transfer_id, block_ids); the worker's start_load_caches

37 + # copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id

38 + # once its copy lands. The id is minted per dispatch, so two loads of the

39 + # same mm_hash in different steps are distinct and a late report from the

40 + # earlier one cannot release the later one's pin.

41 + loads: dict[str, tuple[int, list[int]]] = field(default_factory=dict)

42 +

43 +

44 +@dataclass

45 +class ECCPUWorkerMetadata(ECConnectorWorkerMetadata):

46 + """Per-step worker → scheduler payload for the ECCPUConnector.

47 +

48 + Reports the GPU copies that completed this step: saved mm_hashes become

49 + safe to mark ready, and loaded transfers become safe to unpin once every

50 + participating rank has reported them. Built by

51 + `ECCPUWorker.build_connector_worker_meta`; consumed by

52 + `ECCPUScheduler.update_connector_output`.

53 +

54 + Loads are reported by transfer id rather than mm_hash because every rank

55 + copies the same blocks: `aggregate` concatenates, so an id appears once

56 + per reporting rank and the scheduler can count participants off the list.

57 + """

58 +

59 + completed_saves: list[str] = field(default_factory=list)

60 + completed_loads: list[int] = field(default_factory=list)

61 +

62 + def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetadata:

63 + assert isinstance(other, ECCPUWorkerMetadata)

64 + self.completed_saves.extend(other.completed_saves)

65 + self.completed_loads.extend(other.completed_loads)

66 + return self

69 def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int:

vllm/distributed/ec_transfer/ec_connector/cpu/connector.py

+16

@@ -15,6 +15,7 @@ import torch

15 from vllm.distributed.ec_transfer.ec_connector.base import (

16 ECConnectorBase,

17 ECConnectorRole,

18 + ECConnectorWorkerMetadata,

19 )

20 from vllm.distributed.ec_transfer.ec_connector.cpu.common import (

21 ECCPUConnectorMetadata,

@@ -24,6 +25,7 @@ from vllm.logger import init_logger

25 if TYPE_CHECKING:

26 from vllm.config import VllmConfig

27 from vllm.v1.core.sched.output import SchedulerOutput

28 + from vllm.v1.outputs import ECConnectorOutput

29 from vllm.v1.request import Request

31 logger = init_logger(__name__)

@@ -113,6 +115,20 @@ class ECCPUConnector(ECConnectorBase):

115 self.connector_worker.flush_saves()

116 return None, None

117

118 + def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None:

119 + if self.connector_worker is not None:

120 + return self.connector_worker.build_connector_worker_meta()

121 + return None

122 +

123 + def update_connector_output(self, connector_output: "ECConnectorOutput") -> None:

124 + if self.connector_scheduler is not None:

125 + self.connector_scheduler.update_connector_output(connector_output)

126 +

127 + def has_pending_push_work(self) -> bool:

128 + if self.connector_scheduler is not None:

129 + return self.connector_scheduler.has_pending_push_work()

130 + return False

131 +

132 # Shared.

133 def shutdown(self) -> None:

134 if self.connector_scheduler is not None:

vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py

+63 -31

@@ -11,19 +11,18 @@ from typing import TYPE_CHECKING

12 from vllm.distributed.ec_transfer.ec_connector.cpu.common import (

13 ECCPUConnectorMetadata,

14 + ECCPUWorkerMetadata,

15 create_ec_shared_region,

16 )

17 from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (

18 EmbeddingCache,

19 )

19 -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.step_tracker import (

20 - StepTracker,

21 -)

20 from vllm.logger import init_logger

22 if TYPE_CHECKING:

23 from vllm.config import VllmConfig

24 from vllm.v1.core.sched.output import SchedulerOutput

25 + from vllm.v1.outputs import ECConnectorOutput

26 from vllm.v1.request import Request

28 logger = init_logger(__name__)

@@ -42,16 +41,27 @@ class ECCPUScheduler:

41 # Block allocator + LRU eviction policy for the shared region.

42 self._cache = EmbeddingCache(self._region.num_blocks)

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 -

44 # mm_hash → block IDs allocated this step for GPU→mmap saves.

45 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]] = {}

46 + # mm_hash → (transfer_id, block IDs) to load from mmap→GPU this step.

47 + self._pending_loads: dict[str, tuple[int, list[int]]] = {}

48 +

49 + # Dispatched loads awaiting completion reports, keyed by transfer id:

50 + # transfer_id → (mm_hash, reports still outstanding). The pin taken at

51 + # dispatch is released only once the count reaches zero.

52 + self._load_acks: dict[int, tuple[str, int]] = {}

53 + self._next_transfer_id = 0

54 +

55 + pc = vllm_config.parallel_config

56 + # Reports to expect per load. Only pipeline stage 0 runs the EC

57 + # connector, so the participants are the tp × pcp ranks of that stage.

58 + # Other executors deliver a single rank's output, and expecting reports

59 + # that cannot arrive would hold the pin forever.

60 + self._expected_load_reports = (

61 + pc.tensor_parallel_size * pc.prefill_context_parallel_size

62 + if pc.distributed_executor_backend == "mp"

63 + else 1

64 + )

66 def has_cache_item(self, identifier: str) -> bool:

67 if not self._is_consumer:

@@ -72,30 +82,19 @@ class ECCPUScheduler:

82 entry = self._cache.alloc(mm_hash, feature.mm_position.length)

83 if entry is not None:

84 self._pending_saves[mm_hash] = list(entry.block_ids)

75 - self._ready_tracker.add(mm_hash, request.request_id)

86 if self._is_consumer and mm_hash not in self._pending_loads:

87 entry = self._cache.get(mm_hash)

88 if entry is not None and entry.ready:

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

90 + transfer_id = self._next_transfer_id

91 + self._next_transfer_id += 1

92 + self._pending_loads[mm_hash] = (transfer_id, list(entry.block_ids))

93 + self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports)

95 def build_connector_meta(

96 self, scheduler_output: "SchedulerOutput"

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

98 meta = ECCPUConnectorMetadata()

99 if self._is_producer:

100 meta.saves = self._pending_saves

@@ -105,13 +104,46 @@ class ECCPUScheduler:

104 self._pending_loads = {}

105 return meta

106

107 + def update_connector_output(self, connector_output: "ECConnectorOutput") -> None:

108 + """Apply the worker's memcpy-completion report to the cache.

109 +

110 + Completed saves become safe to mark ready. A load is unpinned once

111 + every participating rank has reported its transfer id; reports for a

112 + transfer that has already been released, or for one this scheduler

113 + never dispatched, are ignored.

114 + """

115 + meta = connector_output.ec_connector_worker_meta

116 + if not isinstance(meta, ECCPUWorkerMetadata):

117 + return

118 + for mm_hash in meta.completed_saves:

119 + entry = self._cache.get(mm_hash)

120 + if entry is not None and not entry.ready:

121 + self._cache.mark_ready(mm_hash)

122 + for transfer_id in meta.completed_loads:

123 + pending = self._load_acks.get(transfer_id)

124 + if pending is None:

125 + continue

126 + mm_hash, outstanding = pending

127 + if outstanding > 1:

128 + self._load_acks[transfer_id] = (mm_hash, outstanding - 1)

129 + continue

130 + # Drop the entry before unpinning so a replayed report is treated

131 + # as stale rather than releasing the pin a second time.

132 + del self._load_acks[transfer_id]

133 + self._cache.unpin(mm_hash)

134 +

135 + def has_pending_push_work(self) -> bool:

136 + """True while any dispatched save or load has not been confirmed done.

137 +

138 + Keeps the engine stepping so the worker's completion reports are polled

139 + even when no requests are otherwise runnable.

140 + """

141 + return self._cache.has_held_entries()

142 +

143 def shutdown(self) -> None:

109 - # drain_all() covers both entries still in _current (never

110 - # consumed by build_connector_meta) and entries in slots.

144 + self._pending_saves.clear()

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

146 + self._load_acks.clear()

147

148 self._is_producer = False

149 self._is_consumer = False

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

+13 -1

@@ -67,6 +67,15 @@ class EmbeddingCache:

67 with self._lock:

68 return self._entries.get(key)

70 + def has_held_entries(self) -> bool:

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

72 +

73 + Held entries are exactly those absent from the eviction free list, so

74 + they represent saves awaiting mark_ready or loads awaiting unpin.

75 + """

76 + with self._lock:

77 + return len(self._entries) != len(self._entries_free_list)

78 +

79 def alloc(self, key: str, n_blocks: int) -> CacheEntry | None:

80 """Allocate *n_blocks* for *key*, evicting as needed.

@@ -84,7 +93,10 @@ class EmbeddingCache:

93 if len(self._free_blocks) + self._evictable_block_count < n_blocks:

94 return None

95 self._evict_until(n_blocks)

87 - block_ids = tuple(self._free_blocks.pop() for _ in range(n_blocks))

96 + # Canonical ascending order: keeps a freed entry's blocks contiguous

97 + # in the free list when it is evicted, and lets consumers detect

98 + # which of an entry's blocks are adjacent.

99 + block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks)))

100 entry = CacheEntry(block_ids)

101 self._entries[key] = entry

102 return entry

vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/step_tracker.py deleted

-115

@@ -1,115 +0,0 @@

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

+266 -50

@@ -7,13 +7,19 @@ per-step connector metadata (`ECCPUConnectorMetadata`) to decide which

7 blocks to copy in each direction.

8 """

10 +from collections import deque

11 +from collections.abc import Sequence

12 +from dataclasses import dataclass

13 +from itertools import chain

14 from typing import TYPE_CHECKING

16 +import numpy as np

17 import torch

19 from vllm._custom_ops import swap_blocks_batch

20 from vllm.distributed.ec_transfer.ec_connector.cpu.common import (

21 ECCPUConnectorMetadata,

22 + ECCPUWorkerMetadata,

23 create_ec_shared_region,

24 )

25 from vllm.distributed.ec_transfer.ec_connector.cpu.worker.descriptor_buffers import (

@@ -34,6 +40,50 @@ if TYPE_CHECKING:

40 logger = init_logger(__name__)

43 +def _coalesce_runs(

44 + block_ids: Sequence[int],

45 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]:

46 + """Split block ids into runs that are consecutive in the region.

47 +

48 + Returns `(slots, first_blocks, num_blocks)`, one element per run, where

49 + `slots` is the run's offset in the flat block sequence. Per-descriptor

50 + overhead dominates these copies, so describing a run with one descriptor

51 + rather than one per block is worth an order of magnitude.

52 +

53 + Both callers derive the non-region side of a run from its slot, so that

54 + side advances in lockstep with the region and only the region has to be

55 + split. `EmbeddingCache` hands out ascending ids, which is what makes runs

56 + findable at all.

57 + """

58 +

59 + ids = np.fromiter(block_ids, dtype=np.int64, count=len(block_ids))

60 + breaks = np.flatnonzero(np.diff(ids) != 1) + 1

61 + slots = np.concatenate(([0], breaks))

62 + ends = np.concatenate((breaks, [ids.size]))

63 + return slots, ids[slots], ends - slots

64 +

65 +

66 +@dataclass

67 +class Transfer:

68 + """A batched copy in flight on a dedicated GPU stream.

69 +

70 + Its descriptor buffers and stream are held until `end_event` fires; only

71 + then are the buffers and stream safe to recycle and the `completions` safe

72 + to report. `start_event`/`end_event` bracket the copy so its elapsed time

73 + can be measured on completion.

74 +

75 + `completions` is what the scheduler is told about this copy: mm_hashes for

76 + saves, transfer ids for loads.

77 + """

78 +

79 + start_event: torch.Event

80 + end_event: torch.Event

81 + completions: list[str] | list[int]

82 + bufs: DescriptorBuffers

83 + stream: torch.Stream

84 + num_bytes: int

85 +

86 +

87 class ECCPUWorker:

88 """Worker-side delegate for the ECCPUConnector.

@@ -43,8 +93,18 @@ class ECCPUWorker:

93 batched call in `flush_saves`.

94 - Consumer role: copies `mmap[block_ids]` → `encoder_cache[mm_hash]`

95 for all entries in `metadata.loads` via a single `swap_blocks_batch`

46 - call on the load stream.

96 + call on a dedicated stream.

97 - On `ec_both` nodes both paths run back-to-back in a single step.

98 +

99 + Every batched copy runs on a stream drawn from a pool so it overlaps the

100 + compute stream, bracketed by a start/end event pair. The stream, events,

101 + and descriptor buffers are recycled once the end event fires.

102 +

103 + Both paths move memory between the two streams, so each buffer is registered

104 + against the stream that is not its own: save sources against the save stream,

105 + load destinations against the compute stream. Without that, the caching

106 + allocator would be free to reissue memory a copy or a queued read still

107 + needs, since it only tracks the stream a buffer was allocated on.

108 """

109

110 def __init__(self, vllm_config: "VllmConfig") -> None:

@@ -62,15 +122,72 @@ class ECCPUWorker:

122 get_tensor_model_parallel_rank() == 0 and get_pcp_group().rank_in_group == 0

123 )

124

65 - # Dedicated stream for async mmap→GPU loads (overlaps with compute).

66 - self._load_stream = current_platform.Stream()

67 -

125 # Descriptor buffer pool (recycled across steps, shared by both paths).

126 self._buf_pool = DescriptorBufferPool()

127

128 # Active save buffer being filled during save_caches calls this step.

129 self._save_bufs: DescriptorBuffers | None = None

130 self._save_count: int = 0

131 + # Bytes accumulated into the active save buffer this step, for logging.

132 + self._save_bytes: int = 0

133 + # mm_hashes accumulated into the active save buffer this step.

134 + self._save_mm_hashes: list[str] = []

135 + # Stream the active save buffer's copy will run on, held from the moment

136 + # the buffer opens so `save_caches` can register its sources against it.

137 + self._save_stream: torch.Stream | None = None

138 +

139 + # Batched copies whose end event has not yet fired. Buffers and stream

140 + # stay held until the event fires; the completions are reported to the

141 + # scheduler only then. Separate per direction because saves and loads

142 + # can be in flight concurrently.

143 + self._inflight_saves: deque[Transfer] = deque()

144 + self._inflight_loads: deque[Transfer] = deque()

145 +

146 + # Streams and timing events recycled across steps to avoid per-step

147 + # create/destroy churn. Shared by both directions.

148 + self._stream_pool: list[torch.Stream] = []

149 + self._event_pool: list[torch.Event] = []

150 +

151 + def _acquire_stream(self) -> torch.Stream:

152 + if self._stream_pool:

153 + return self._stream_pool.pop()

154 + return current_platform.Stream()

155 +

156 + def _acquire_event(self) -> torch.Event:

157 + return (

158 + self._event_pool.pop()

159 + if self._event_pool

160 + else torch.Event(enable_timing=True)

161 + )

162 +

163 + def _collect_finished(self, inflight: deque[Transfer], direction: str) -> list:

164 + """Pop transfers whose end event has fired, recycle their stream,

165 + events, and buffers, and return their completions.

166 +

167 + The front check is conservative: a transfer is reported only once its

168 + own end event fires, so completions are always genuinely done, and each

169 + is reported exactly once because the transfer is popped as it is

170 + reported. A later transfer that happens to finish first simply waits

171 + behind the front and is reported on a subsequent poll.

172 + """

173 + done: list = []

174 + while inflight and inflight[0].end_event.query():

175 + transfer = inflight.popleft()

176 + done.extend(transfer.completions)

177 + elapsed_ms = transfer.start_event.elapsed_time(transfer.end_event)

178 + logger.debug(

179 + "EC %s: %d entr%s (%d bytes) took %.3f ms",

180 + direction,

181 + len(transfer.completions),

182 + "y" if len(transfer.completions) == 1 else "ies",

183 + transfer.num_bytes,

184 + elapsed_ms,

185 + )

186 + self._buf_pool.release(transfer.bufs)

187 + self._stream_pool.append(transfer.stream)

188 + self._event_pool.append(transfer.start_event)

189 + self._event_pool.append(transfer.end_event)

190 + return done

191

192 def save_caches(

193 self,

@@ -78,7 +195,14 @@ class ECCPUWorker:

195 mm_hash: str,

196 connector_metadata: ECCPUConnectorMetadata,

197 ) -> None:

81 - """Fill descriptor buffers directly for batched flush."""

198 + """Fill descriptor buffers directly for batched flush.

199 +

200 + Registers `encoder_cache[mm_hash]` against the stream the copy will run

201 + on, so the caching allocator will not hand that memory to another

202 + allocation before the copy has read it. The descriptor buffers hold raw

203 + addresses, and the caller is free to evict the entry as soon as this

204 + step's saves are dispatched.

205 + """

206 if not self._is_save_rank:

207 return

208 block_ids = connector_metadata.saves.get(mm_hash)

@@ -99,43 +223,95 @@ class ECCPUWorker:

223 if self._save_bufs is None:

224 total = sum(len(v) for v in connector_metadata.saves.values())

225 self._save_bufs = self._buf_pool.acquire(total)

226 + self._save_stream = self._acquire_stream()

227 + self._save_mm_hashes = []

228

229 assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel()

230

231 bufs = self._save_bufs

106 - src_base = src.view(-1).view(torch.uint8).data_ptr()

232 + assert self._save_stream is not None

233 + src.record_stream(self._save_stream)

234 + # Device addresses are unsigned and can have the high bit set (XPU USM),

235 + # so they are added up in uint64; signed arithmetic rejects them.

236 + src_base = np.uint64(src.view(-1).view(torch.uint8).data_ptr())

237 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 - bufs.set_ptrs(idx, src_base + start, dst_base + block_idx * block_size)

113 - bufs.sizes[idx] = min(block_size, total_bytes - start)

114 - idx += 1

238 + # Entries are sized by placeholder count, which can exceed the number

239 + # of embeddings, so blocks the output never reaches get no descriptor.

240 + n_used = -(-total_bytes // block_size)

241 + if n_used:

242 + slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used])

243 + run_bytes = num_blocks * block_size

244 + # The output can stop partway into the final run.

245 + run_bytes[-1] = total_bytes - slots[-1] * block_size

246 + bufs.add_copies(

247 + self._save_count,

248 + src_base + slots.astype(np.uint64) * np.uint64(block_size),

249 + dst_base + first_blocks * block_size,

250 + run_bytes,

251 + )

252 + self._save_count += slots.size

253

116 - self._save_count = idx

254 + self._save_bytes += total_bytes

255 + self._save_mm_hashes.append(mm_hash)

256

257 def flush_saves(self) -> None:

119 - """Flush all accumulated saves in a single swap_blocks_batch call."""

258 + """Flush all accumulated saves in a single swap_blocks_batch call.

259 +

260 + Runs the copy on the stream the batch acquired when it opened, gated

261 + behind the compute stream that produced the encoder outputs, brackets it

262 + with start/end events, and enqueues the batch as in-flight. The stream,

263 + events, and descriptor buffers are recycled once the end event fires (see

264 + `_collect_finished`), which is also when the saved mm_hashes become

265 + safe to mark ready.

266 + """

267 if self._save_count == 0:

268 return

269

270 bufs = self._save_bufs

124 - assert bufs is not None

271 + stream = self._save_stream

272 + assert bufs is not None and stream is not None

273 src_ptrs, dst_ptrs, sizes = bufs.src_ptrs, bufs.dst_ptrs, bufs.sizes

274 n = self._save_count

127 - swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n])

275 + num_bytes = self._save_bytes

276 +

277 + # Gate the GPU→CPU copy behind the compute stream: it reads the encoder

278 + # outputs the model just produced there.

279 + stream.wait_stream(current_platform.current_stream())

280 + start_event = self._acquire_event()

281 + end_event = self._acquire_event()

282 + with current_platform.stream(stream):

283 + start_event.record(stream)

284 + swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n])

285 + end_event.record(stream)

286 +

287 + self._inflight_saves.append(

288 + Transfer(

289 + start_event=start_event,

290 + end_event=end_event,

291 + completions=self._save_mm_hashes,

292 + bufs=bufs,

293 + stream=stream,

294 + num_bytes=num_bytes,

295 + )

296 + )

297

129 - self._buf_pool.release(bufs)

298 self._save_bufs = None

299 + self._save_stream = None

300 self._save_count = 0

301 + self._save_bytes = 0

302 + self._save_mm_hashes = []

303

304 def start_load_caches(

305 self,

306 encoder_cache: dict[str, torch.Tensor],

307 connector_metadata: ECCPUConnectorMetadata,

308 ) -> None:

138 - """Consumer path: single batched copy of all loads from mmap→GPU."""

309 + """Consumer path: single batched copy of all loads from mmap→GPU.

310 +

311 + The destination is registered against the compute stream that consumes

312 + it, so the caching allocator will not hand that memory to a later load

313 + while reads of it are still queued there.

314 + """

315 if not connector_metadata.loads:

316 return

317

@@ -145,59 +321,99 @@ class ECCPUWorker:

321 device_type = current_platform.device_type

322 src_base = blocks.data_ptr()

323

148 - # Pre-filter: only hashes not already in encoder_cache.

149 - load_items = {

150 - h: idxs

151 - for h, idxs in connector_metadata.loads.items()

152 - if h not in encoder_cache

153 - }

154 - if not load_items:

155 - return

156 -

157 - total_blocks = sum(len(idxs) for idxs in load_items.values())

324 + # Copy every dispatched load. The scheduler routes an input here only

325 + # after missing the GPU encoder cache, so a resident hit is not expected;

326 + # copying unconditionally means every participating rank reports each

327 + # transfer exactly once, which is what the scheduler counts to release

328 + # the pin.

329 + # Ordered by first block so entries that neighbour each other in the

330 + # region land next to each other in the destination and coalesce into

331 + # one descriptor. Every loop below consumes this same order.

332 + load_items = sorted(

333 + connector_metadata.loads.items(), key=lambda kv: kv[1][1][0]

334 + )

335 + total_blocks = sum(len(block_ids) for _, (_, block_ids) in load_items)

336

159 - with current_platform.stream(self._load_stream):

337 + stream = self._acquire_stream()

338 + compute_stream = current_platform.current_stream()

339 + start_event = self._acquire_event()

340 + end_event = self._acquire_event()

341 + with current_platform.stream(stream):

342 # Single contiguous destination buffer for all loads.

343 dst_buf = torch.empty(

344 total_blocks, block_size, dtype=torch.int8, device=device_type

345 )

164 - dst_buf_base = dst_buf.data_ptr()

346 + dst_buf.record_stream(compute_stream)

347 + # See save_caches: device addresses are added up in uint64.

348 + dst_buf_base = np.uint64(dst_buf.data_ptr())

349

350 bufs = self._buf_pool.acquire(total_blocks)

167 - src_ptrs = bufs.src_ptrs[:total_blocks]

168 - dst_ptrs = bufs.dst_ptrs[:total_blocks]

169 - sizes = bufs.sizes[:total_blocks]

170 - sizes[:] = block_size

171 -

172 - op_idx = 0

173 - for block_ids in load_items.values():

174 - for block_idx in block_ids:

175 - bufs.set_ptrs(

176 - op_idx,

177 - src_base + block_idx * block_size,

178 - dst_buf_base + op_idx * block_size,

179 - )

180 - op_idx += 1

351 + slots, first_blocks, num_blocks = _coalesce_runs(

352 + list(chain.from_iterable(ids for _, (_, ids) in load_items))

353 + )

354 + bufs.add_copies(

355 + 0,

356 + src_base + first_blocks * block_size,

357 + dst_buf_base + slots.astype(np.uint64) * np.uint64(block_size),

358 + num_blocks * block_size,

359 + )

360 + op_idx = slots.size

361 + src_ptrs = bufs.src_ptrs[:op_idx]

362 + dst_ptrs = bufs.dst_ptrs[:op_idx]

363 + sizes = bufs.sizes[:op_idx]

364

365 + start_event.record(stream)

366 swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any=True)

183 -

184 - self._buf_pool.release(bufs)

367 + end_event.record(stream)

368 + self._inflight_loads.append(

369 + Transfer(

370 + start_event=start_event,

371 + end_event=end_event,

372 + completions=[tid for _, (tid, _) in load_items],

373 + bufs=bufs,

374 + stream=stream,

375 + num_bytes=total_blocks * block_size,

376 + )

377 + )

378

379 # Slice contiguous buffer into per-hash views.

380 offset = 0

188 - for mm_hash, block_ids in load_items.items():

381 + for mm_hash, (_, block_ids) in load_items:

382 n = len(block_ids)

383 encoder_cache[mm_hash] = (

384 dst_buf[offset : offset + n].view(dtype).reshape(n, -1)

385 )

386 offset += n

387

195 - current_platform.current_stream().wait_stream(self._load_stream)

388 + current_platform.current_stream().wait_stream(stream)

389 +

390 + def build_connector_worker_meta(self) -> ECCPUWorkerMetadata | None:

391 + """Report the GPU copies that completed on this rank this step: saved

392 + mm_hashes and loaded transfer ids.

393 +

394 + Returns None when nothing finished, so the scheduler sees no payload.

395 + """

396 + completed_saves = self._collect_finished(self._inflight_saves, "save")

397 + completed_loads = self._collect_finished(self._inflight_loads, "load")

398 + if not completed_saves and not completed_loads:

399 + return None

400 + return ECCPUWorkerMetadata(

401 + completed_saves=completed_saves,

402 + completed_loads=completed_loads,

403 + )

404

405 def shutdown(self) -> None:

198 - self._load_stream.synchronize()

406 + for transfer in (*self._inflight_saves, *self._inflight_loads):

407 + transfer.end_event.synchronize()

408 self._save_bufs = None

409 + self._save_stream = None

410 self._save_count = 0

411 + self._save_bytes = 0

412 + self._save_mm_hashes = []

413 + self._inflight_saves.clear()

414 + self._inflight_loads.clear()

415 + self._stream_pool.clear()

416 + self._event_pool.clear()

417 try:

418 self._region.cleanup()

419 except Exception:

vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py

+36 -15

@@ -22,24 +22,43 @@ class DescriptorBuffers(NamedTuple):

22 src_ptrs: torch.Tensor

23 dst_ptrs: torch.Tensor

24 sizes: torch.Tensor

25 - # Numpy aliases of src_ptrs/dst_ptrs, written via set_ptrs().

25 + # Numpy aliases of the three tensors, written via add_copies().

26 src_np: np.ndarray

27 dst_np: np.ndarray

28 + sizes_np: np.ndarray

29 - def set_ptrs(self, idx: int, src: int, dst: int) -> None:

30 - """Record the source and destination address of descriptor *idx*.

31 -

32 - TODO(torch>=2.14): drop this indirection and assign the tensors

33 - directly once the minimum supported torch is 2.14. The numpy detour

34 - exists only because torch's setitem unpacks the value as a signed

35 - long long before 2.14 (pytorch#191458), rejecting XPU USM pointers

36 - >= 2**63, while the two's-complement rewrite those versions accept is

37 - in turn rejected by uint64 from 2.14 on. Numpy casts against the

38 - array dtype and so works on either. `sizes` holds byte counts and

39 - needs no such care.

30 + def add_copies(

31 + self,

32 + first: int,

33 + sources: np.ndarray,

34 + destinations: np.ndarray,

35 + byte_counts: np.ndarray,

36 + ) -> None:

37 + """Store a batch of copies in these buffers, the first at *first*.

38 +

39 + One copy is one source address, one destination address, and a number

40 + of bytes. `swap_blocks_batch` performs every copy held in the buffers

41 + in a single call, so a caller fills in all of a step's copies before

42 + asking for any of them to run, and passes *first* to say where its own

43 + copies begin.

44 +

45 + Whole batches go in at once because storing a single number into a

46 + torch tensor costs microseconds. Adding copies one at a time can take

47 + longer than performing them.

48 +

49 + TODO(torch>=2.14): write the tensors directly and drop these numpy

50 + aliases once the minimum supported torch is 2.14. They exist only

51 + because torch's setitem unpacks the value as a signed long long before

52 + 2.14 (pytorch#191458), rejecting XPU USM addresses >= 2**63, while the

53 + two's-complement rewrite those versions accept is in turn rejected by

54 + uint64 from 2.14 on. Numpy casts against the array's own type and so

55 + works on either. Byte counts are small enough to need none of this

56 + care, and use their alias only to match the addresses.

57 """

41 - self.src_np[idx] = src

42 - self.dst_np[idx] = dst

58 + last = first + len(sources)

59 + self.src_np[first:last] = sources

60 + self.dst_np[first:last] = destinations

61 + self.sizes_np[first:last] = byte_counts

64 class DescriptorBufferPool:

@@ -62,7 +81,9 @@ class DescriptorBufferPool:

81 if bufs.src_ptrs.numel() >= n:

82 return bufs

83 src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3))

65 - return DescriptorBuffers(src, dst, sizes, src.numpy(), dst.numpy())

84 + return DescriptorBuffers(

85 + src, dst, sizes, src.numpy(), dst.numpy(), sizes.numpy()

86 + )

88 def release(self, bufs: DescriptorBuffers) -> None:

89 """Return a buffer triple to the pool for reuse."""