iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
omerpaz95
2026-08-27 03:41:46 +00:00
committed by GitHub
co-authored by Claude Opus 5
parent 94d96e2446
commit f8d38fbc87
15 changed files with 1054 additions and 650 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/ec_connector/__init__.py
View File
tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py
+30 -1
View File
@@ -21,13 +21,15 @@ def _cache(num_blocks: int = 8) -> EmbeddingCache:
# ── alloc ────────────────────────────────────────────────────────────────────
def test_alloc_returns_unique_block_ids():
def test_alloc_returns_unique_ascending_block_ids():
cache = _cache()
entry = cache.alloc("a", 3)
assert entry is not None
assert len(entry.block_ids) == 3
assert len(set(entry.block_ids)) == 3
assert all(0 <= i < 8 for i in entry.block_ids)
# Ascending order lets consumers detect which blocks are adjacent.
assert list(entry.block_ids) == sorted(entry.block_ids)
def test_alloc_returns_tuple():
@@ -161,6 +163,33 @@ def test_unpin_asserts_on_unpinned():
cache.unpin("a")
# ── has_held_entries ─────────────────────────────────────────────────────────
def test_has_held_entries_empty_cache():
cache = _cache()
assert cache.has_held_entries() is False
def test_has_held_entries_true_while_not_ready():
cache = _cache()
cache.alloc("a", 2) # not ready → held
assert cache.has_held_entries() is True
cache.mark_ready("a") # ready + unpinned → settled
assert cache.has_held_entries() is False
def test_has_held_entries_true_while_pinned():
cache = _cache()
cache.alloc("a", 2)
cache.mark_ready("a")
assert cache.has_held_entries() is False
cache.pin("a") # pinned → held
assert cache.has_held_entries() is True
cache.unpin("a")
assert cache.has_held_entries() is False
# ── eviction ─────────────────────────────────────────────────────────────────
tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py
+165 -183
View File
@@ -3,6 +3,11 @@
import uuid
import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
from vllm.config.ec_transfer import ECRole
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUWorkerMetadata,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
)
@@ -37,13 +42,26 @@ class _Request:
self.request_id = request_id
class _WorkerOutput:
"""Stand-in for ECConnectorOutput carrying a completion report.
`saves` are mm_hashes; `loads` are transfer ids. One entry per reporting
rank, so a list holding the same id twice is two ranks reporting it.
"""
def __init__(self, *, saves=None, loads=None):
self.ec_connector_worker_meta = ECCPUWorkerMetadata(
completed_saves=list(saves or []),
completed_loads=list(loads or []),
)
def _make_scheduler(
monkeypatch,
num_blocks=_N,
max_concurrent_batches=1,
*,
is_producer=True,
is_consumer=True,
ec_role: ECRole = "ec_both",
tensor_parallel_size=1,
) -> ECCPUScheduler:
region = ECSharedRegion(
engine_id=str(uuid.uuid4()),
@@ -52,27 +70,22 @@ def _make_scheduler(
)
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)
_is_prod = is_producer
_is_cons = is_consumer
_mcb = max_concurrent_batches
class _EC:
is_ec_producer = _is_prod
is_ec_consumer = _is_cons
class _Cfg:
ec_transfer_config = _EC()
@property
def max_concurrent_batches(self):
return _mcb
return ECCPUScheduler(_Cfg())
return ECCPUScheduler(
create_ec_vllm_config(
ec_role=ec_role,
tensor_parallel_size=tensor_parallel_size,
)
)
class _SchedulerOutput:
def __init__(self, finished_req_ids=None):
self.finished_req_ids = finished_req_ids or set()
def _load_ids(meta) -> list[int]:
"""Transfer ids the scheduler dispatched in this step's metadata."""
return [transfer_id for transfer_id, _ in meta.loads.values()]
def _load_blocks(meta, mm_hash: str) -> list[int]:
"""Block ids the scheduler dispatched for one load."""
return meta.loads[mm_hash][1]
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):
def test_offload_reuse_cycle(monkeypatch):
s = _make_scheduler(monkeypatch)
req = _Request([_Feature("h1", length=1)])
# Two blocks, so the reload also pins down that every allocated block id
# comes back in meta.loads, not just the first.
req = _Request([_Feature("h1", length=2)])
# Step A: first sight — allocate and save.
assert s.has_cache_item("h1") is False
@@ -92,24 +107,24 @@ def test_offload_reuse_cycle(monkeypatch):
meta_a = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta_a.saves
assert meta_a.loads == {}
# Not ready yet (max_concurrent_batches=1 means ready next step).
# Not ready until the worker reports the save memcpy complete.
assert s.has_cache_item("h1") is False
# Step B: _step_readiness marks h1 ready (no new allocs).
s.build_connector_meta(scheduler_output=None)
# Worker reports the save copy finished → entry becomes ready.
s.update_connector_output(_WorkerOutput(saves=["h1"]))
assert s.has_cache_item("h1") is True
# Step C: reload now works — entry is ready at update_state_after_alloc time.
s.update_state_after_alloc(req, 0)
meta_c = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta_c.loads
assert meta_c.loads["h1"] == meta_a.saves["h1"]
assert _load_blocks(meta_c, "h1") == meta_a.saves["h1"]
s.shutdown()
def test_has_cache_item_false_when_not_consumer(monkeypatch):
s = _make_scheduler(monkeypatch, is_consumer=False)
s = _make_scheduler(monkeypatch, ec_role="ec_producer")
assert s.has_cache_item("anything") is False
s.shutdown()
@@ -123,10 +138,10 @@ def test_connector_keys_on_identifier_not_mm_hash(monkeypatch):
meta = s.build_connector_meta(scheduler_output=None)
assert "ENC_KEY" in meta.saves
assert "PROC_KEY" not in meta.saves
# Not ready yet (1 step delay).
# Not ready until the completion report.
assert s.has_cache_item("ENC_KEY") is False
# After another step it becomes ready — keyed on identifier.
s.build_connector_meta(scheduler_output=None)
# After the report it becomes ready — keyed on identifier.
s.update_connector_output(_WorkerOutput(saves=["ENC_KEY"]))
assert s.has_cache_item("ENC_KEY") is True
assert s.has_cache_item("PROC_KEY") is False
s.shutdown()
@@ -136,13 +151,11 @@ def test_cpu_region_fifo_eviction(monkeypatch):
# Region holds exactly 2 one-block encodings.
s = _make_scheduler(monkeypatch, num_blocks=2)
# Fill with A and B, marking each ready.
# Fill with A and B, marking each ready via a completion report.
for h in ("A", "B"):
s.update_state_after_alloc(_Request([_Feature(h, length=1)]), 0)
s.build_connector_meta(scheduler_output=None)
# After the step above, A's step was called in B's build. A is ready.
# B was submitted in the last build, so one more step makes it ready.
s.build_connector_meta(scheduler_output=None) # marks B ready
s.update_connector_output(_WorkerOutput(saves=[h]))
assert s.has_cache_item("A") is True
assert s.has_cache_item("B") is True
@@ -161,8 +174,8 @@ def test_multiple_mm_items_per_request(monkeypatch):
meta = s.build_connector_meta(scheduler_output=None)
assert set(meta.saves) == {"h1", "h2"}
# Step delay: mark both ready.
s.build_connector_meta(scheduler_output=None)
# Worker reports both saves complete.
s.update_connector_output(_WorkerOutput(saves=["h1", "h2"]))
assert s.has_cache_item("h1") is True
assert s.has_cache_item("h2") is True
@@ -174,26 +187,6 @@ def test_multiple_mm_items_per_request(monkeypatch):
s.shutdown()
def test_load_returns_correct_block_ids(monkeypatch):
"""meta.loads must contain the same block IDs that meta.saves
allocated — verified through the public API only."""
s = _make_scheduler(monkeypatch)
req = _Request([_Feature("a", length=2)])
s.update_state_after_alloc(req, 0)
meta_save = s.build_connector_meta(scheduler_output=None)
assert "a" in meta_save.saves
# Readiness delay.
s.build_connector_meta(scheduler_output=None)
# Reload.
s.update_state_after_alloc(req, 0)
meta_load = s.build_connector_meta(scheduler_output=None)
assert meta_load.loads["a"] == meta_save.saves["a"]
s.shutdown()
def test_loads_only_serves_hashes_touched_this_step(monkeypatch):
s = _make_scheduler(monkeypatch)
_seed_cached(s, "a", n_blocks=2)
@@ -228,10 +221,10 @@ def test_load_not_emitted_for_uncached_entry(monkeypatch):
s.shutdown()
def test_delayed_unpin_protects_blocks_during_worker_read(monkeypatch):
"""Blocks being loaded must not be evictable until the delayed unpin
fires. Observable: a save that needs eviction cannot reclaim a
still-pinned loaded entry."""
def test_load_pin_protects_blocks_until_completion(monkeypatch):
"""Loaded blocks stay pinned until the worker reports the load memcpy
complete. Observable: a save that needs eviction cannot reclaim a
still-pinned loaded entry, but can once the completion report arrives."""
# 2 blocks, both occupied by "a" (ready).
s = _make_scheduler(monkeypatch, num_blocks=2)
_seed_cached(s, "a", n_blocks=2)
@@ -240,94 +233,79 @@ def test_delayed_unpin_protects_blocks_during_worker_read(monkeypatch):
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "a" in meta.loads
(transfer_id,) = _load_ids(meta)
# Try to save "b" (needs 2 blocks, but "a" is pinned) — must fail.
s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
meta2 = s.build_connector_meta(scheduler_output=None)
assert "b" not in meta2.saves # can't evict pinned "a"
# After the delay, "a" unpins. Now "b" can evict it.
# The one participant reports the load complete → "a" unpins. Now "b" can
# evict it.
s.update_connector_output(_WorkerOutput(loads=[transfer_id]))
s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
meta3 = s.build_connector_meta(scheduler_output=None)
assert "b" in meta3.saves
s.shutdown()
def test_delayed_unpin_depth_2(monkeypatch):
"""With max_concurrent_batches=2, loaded blocks stay pinned for 2 steps
after build. Observable via eviction failure."""
s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=2)
_seed_cached(s, "a", n_blocks=2)
# Load "a" → pinned.
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
s.build_connector_meta(scheduler_output=None)
# Step 1 after load: still pinned — can't evict for "b".
s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
meta1 = s.build_connector_meta(scheduler_output=None)
assert "b" not in meta1.saves
# Step 2 after load: still pinned.
s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
meta2 = s.build_connector_meta(scheduler_output=None)
assert "b" not in meta2.saves
# Step 3: unpin fires, now eviction works.
s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
meta3 = s.build_connector_meta(scheduler_output=None)
assert "b" in meta3.saves
s.shutdown()
def _can_evict(s: ECCPUScheduler, mm_hash: str, n_blocks: int) -> bool:
"""Whether a new save of `mm_hash` can claim `n_blocks`, i.e. whether the
blocks currently held by other entries are reclaimable."""
s.update_state_after_alloc(_Request([_Feature(mm_hash, length=n_blocks)]), 0)
meta = s.build_connector_meta(scheduler_output=None)
return mm_hash in meta.saves
def test_shutdown_unpins_pending_loads(monkeypatch):
"""Shutdown must unpin entries still in _pending_loads (not yet drained
to the unpin deque). Observable: after shutdown, the entry is evictable
(alloc can reclaim its blocks)."""
s = _make_scheduler(monkeypatch, num_blocks=2)
def test_load_needs_every_participant_before_unpin(monkeypatch):
"""With two participating ranks, one report must not release the pin.
This is the multi-rank case: each rank copies the same CPU blocks into its
own GPU memory, so the blocks stay in use until the last rank is done.
"""
s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2)
_seed_cached(s, "a", n_blocks=2)
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
# Don't call build_connector_meta — "a" stays in _pending_loads.
(transfer_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))
# First rank reports: still in use by the second, so "a" must hold.
s.update_connector_output(_WorkerOutput(loads=[transfer_id]))
assert _can_evict(s, "b", 2) is False
assert s.has_pending_push_work() is True
# Second rank reports: now it is genuinely free.
s.update_connector_output(_WorkerOutput(loads=[transfer_id]))
assert s.has_pending_push_work() is False
assert _can_evict(s, "c", 2) is True
s.shutdown()
# After shutdown, scheduler reports no items (consumer role is off).
assert s.has_cache_item("a") is False
def test_late_report_does_not_release_a_later_load_of_same_hash(monkeypatch):
"""A straggling report from one dispatch must not release the pin taken by
a subsequent dispatch of the same mm_hash.
def test_shutdown_unpins_deferred_unpin_queue(monkeypatch):
"""Shutdown must unpin entries queued in the deferred unpin deque
(already drained from _pending_loads but not yet expired)."""
s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=2)
This is what the transfer id is for: mm_hash identifies a cache entry, not
a transfer, so it cannot tell the two dispatches apart.
"""
s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2)
_seed_cached(s, "a", n_blocks=2)
# Load "a" and build — moves to unpin deque.
# First dispatch, fully reported by both ranks → pin released.
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
s.build_connector_meta(scheduler_output=None)
(first_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))
s.update_connector_output(_WorkerOutput(loads=[first_id, first_id]))
assert s.has_pending_push_work() is False
# "a" is in the deferred unpin deque (depth=2, not yet expired).
# Shutdown must unpin it.
s.shutdown()
assert s.has_cache_item("a") is False
# Second dispatch of the same hash takes a fresh pin.
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
(second_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))
assert second_id != first_id
def test_eviction_skips_entry_pinned_by_pending_load(monkeypatch):
"""A loaded entry still within the unpin delay window must survive
eviction attempts from new saves."""
# 2 blocks: A occupies both, ready.
s = _make_scheduler(monkeypatch, num_blocks=2)
_seed_cached(s, "A", n_blocks=2)
# Load A → pinned.
s.update_state_after_alloc(_Request([_Feature("A")]), 0)
meta_load = s.build_connector_meta(scheduler_output=None)
assert "A" in meta_load.loads
# Try to save C (needs 2 blocks) — can't evict pinned A.
s.update_state_after_alloc(_Request([_Feature("C", length=2)]), 0)
meta_save = s.build_connector_meta(scheduler_output=None)
assert "C" not in meta_save.saves
assert s.has_cache_item("A") is True # survived
# A replayed report for the first dispatch must be ignored entirely.
s.update_connector_output(_WorkerOutput(loads=[first_id, first_id]))
assert s.has_pending_push_work() is True
assert _can_evict(s, "b", 2) is False
s.shutdown()
@@ -350,30 +328,10 @@ def test_region_full_skips_save_and_never_blocks(monkeypatch):
s.shutdown()
def test_step_readiness_depth_2(monkeypatch):
"""With max_concurrent_batches=2, entries need 2 steps to become ready."""
s = _make_scheduler(monkeypatch, max_concurrent_batches=2)
req = _Request([_Feature("h1", length=1)])
s.update_state_after_alloc(req, 0)
meta_0 = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta_0.saves
assert s.has_cache_item("h1") is False
# One step later: still not ready (depth=2).
s.build_connector_meta(scheduler_output=None)
assert s.has_cache_item("h1") is False
# Two steps later: now ready.
s.build_connector_meta(scheduler_output=None)
assert s.has_cache_item("h1") is True
s.shutdown()
def test_producer_only_never_emits_loads(monkeypatch):
"""A producer-only scheduler must never populate meta.loads, even when
entries are ready."""
s = _make_scheduler(monkeypatch, is_consumer=False)
s = _make_scheduler(monkeypatch, ec_role="ec_producer")
req = _Request([_Feature("h1", length=1)])
s.update_state_after_alloc(req, 0)
@@ -381,8 +339,7 @@ def test_producer_only_never_emits_loads(monkeypatch):
assert "h1" in meta.saves
assert meta.loads == {}
# Mark ready.
s.build_connector_meta(scheduler_output=None)
s.update_connector_output(_WorkerOutput(saves=["h1"]))
# Even if we re-encounter the same feature, no load.
s.update_state_after_alloc(req, 0)
@@ -393,7 +350,7 @@ def test_producer_only_never_emits_loads(monkeypatch):
def test_consumer_only_never_emits_saves(monkeypatch):
"""A consumer-only scheduler must never populate meta.saves."""
s = _make_scheduler(monkeypatch, is_producer=False)
s = _make_scheduler(monkeypatch, ec_role="ec_consumer")
_seed_cached(s, "a", n_blocks=1)
req = _Request([_Feature("a", length=1)])
@@ -406,7 +363,7 @@ def test_consumer_only_never_emits_saves(monkeypatch):
def test_consumer_only_has_cache_item(monkeypatch):
"""Consumer-only: has_cache_item works normally for ready entries."""
s = _make_scheduler(monkeypatch, is_producer=False)
s = _make_scheduler(monkeypatch, ec_role="ec_consumer")
assert s.has_cache_item("a") is False
_seed_cached(s, "a", n_blocks=1)
assert s.has_cache_item("a") is True
@@ -431,50 +388,75 @@ def test_save_not_emitted_for_already_cached_entry(monkeypatch):
s.shutdown()
# ── first-finish fast-path (idle engine) ────────────────────────────────────
# ── completion reporting (update_connector_output) ──────────────────────────
def test_first_finish_marks_ready_immediately(monkeypatch):
"""When the originating request finishes, the save entry becomes ready
without waiting for the step-count delay."""
s = _make_scheduler(monkeypatch, max_concurrent_batches=5)
req = _Request([_Feature("h1", length=1)], request_id="req_x")
def test_update_connector_output_ignores_foreign_meta(monkeypatch):
"""A worker output whose payload is not an ECCPUWorkerMetadata is a
no-op — no cache mutation, no crash."""
s = _make_scheduler(monkeypatch)
_seed_cached(s, "a", n_blocks=1)
s.update_state_after_alloc(req, 0)
class _Foreign:
ec_connector_worker_meta = object()
s.update_connector_output(_Foreign()) # must not raise
assert s.has_cache_item("a") is True
s.shutdown()
def test_update_connector_output_ignores_unknown_report(monkeypatch):
"""Reports for a never-seen save hash or a transfer id this scheduler never
dispatched are dropped rather than mutating unrelated state."""
s = _make_scheduler(monkeypatch)
s.update_connector_output(_WorkerOutput(saves=["gone"], loads=[4242]))
assert s.has_cache_item("gone") is False
s.shutdown()
# ── has_pending_push_work ───────────────────────────────────────────────────
def test_has_pending_push_work_tracks_inflight_save(monkeypatch):
"""A dispatched-but-unconfirmed save keeps push work pending until the
completion report marks it ready."""
s = _make_scheduler(monkeypatch)
assert s.has_pending_push_work() is False
s.update_state_after_alloc(_Request([_Feature("h1", length=1)]), 0)
s.build_connector_meta(scheduler_output=None)
assert s.has_cache_item("h1") is False
assert s.has_pending_push_work() is True # save not yet confirmed
# Request finishes — readiness should fire immediately.
s.build_connector_meta(_SchedulerOutput(finished_req_ids={"req_x"}))
assert s.has_cache_item("h1") is True
s.update_connector_output(_WorkerOutput(saves=["h1"]))
assert s.has_pending_push_work() is False
s.shutdown()
def test_first_finish_unpins_immediately(monkeypatch):
"""When the originating request finishes, the load entry is unpinned
without waiting for the step-count delay."""
s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=5)
_seed_cached(s, "a", n_blocks=2)
def test_has_pending_push_work_tracks_inflight_load(monkeypatch):
"""A dispatched-but-unconfirmed load keeps push work pending until the
unpin completion report arrives."""
s = _make_scheduler(monkeypatch)
_seed_cached(s, "a", n_blocks=1)
assert s.has_pending_push_work() is False # ready + unpinned
req = _Request([_Feature("a")], request_id="req_y")
s.update_state_after_alloc(req, 0)
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "a" in meta.loads
assert s.has_pending_push_work() is True # pinned, awaiting unpin
# "a" is pinned. Try to evict — must fail.
s.update_state_after_alloc(
_Request([_Feature("b", length=2)], request_id="req_z"), 0
)
meta2 = s.build_connector_meta(scheduler_output=None)
assert "b" not in meta2.saves
# req_y finishes — unpin fires in build_connector_meta (fast-path).
s.build_connector_meta(_SchedulerOutput(finished_req_ids={"req_y"}))
# Next step: "a" is now evictable, so "b" can be allocated.
s.update_state_after_alloc(
_Request([_Feature("b", length=2)], request_id="req_w"), 0
)
meta3 = s.build_connector_meta(scheduler_output=None)
assert "b" in meta3.saves
s.update_connector_output(_WorkerOutput(loads=_load_ids(meta)))
assert s.has_pending_push_work() is False
s.shutdown()
# ── shutdown ────────────────────────────────────────────────────────────────
def test_shutdown_disables_roles_and_cleans_region(monkeypatch):
"""After shutdown the scheduler serves no items and the region is
cleaned up."""
s = _make_scheduler(monkeypatch, num_blocks=2)
_seed_cached(s, "a", n_blocks=2)
assert s.has_cache_item("a") is True
s.shutdown()
assert s.has_cache_item("a") is False # consumer role disabled
tests/v1/ec_connector/unit/cpu/scheduler/test_step_tracker.py
-162
View File
@@ -1,162 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for StepTracker."""
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.step_tracker import (
StepTracker,
)
def test_basic_step_count_expiry_depth_1():
"""With max_concurrent_batches=1, entry expires after 1 step."""
t = StepTracker(max_concurrent_batches=1)
t.add("h1", "req_a")
# Step 1: entry was added, committed to slot. Deque fills to maxlen=1,
# but no expiry yet (the entry is in the slot just committed).
result = t.step(set())
assert result == []
# Step 2: deque is full, oldest slot (containing h1) expires.
result = t.step(set())
assert result == ["h1"]
def test_basic_step_count_expiry_depth_2():
"""With max_concurrent_batches=2, entry expires after 2 steps."""
t = StepTracker(max_concurrent_batches=2)
t.add("h1", "req_a")
assert t.step(set()) == [] # step 1: committed
assert t.step(set()) == [] # step 2: deque fills but h1 in newer slot
assert t.step(set()) == ["h1"] # step 3: h1's slot is oldest, expires
def test_first_finish_fast_path():
"""Entry is returned immediately when request_id appears in finished."""
t = StepTracker(max_concurrent_batches=3)
t.add("h1", "req_a")
result = t.step({"req_a"})
assert result == ["h1"]
def test_no_double_processing_fast_path_then_expiry():
"""Entry fast-pathed should NOT fire again on deque expiry."""
t = StepTracker(max_concurrent_batches=1)
t.add("h1", "req_a")
# Fast-path fires h1.
result = t.step({"req_a"})
assert result == ["h1"]
# Normal expiry would fire now, but h1 is already processed.
result = t.step(set())
assert result == []
def test_no_double_processing_expiry_then_fast_path():
"""Entry expired via step-count should NOT fire again via fast-path."""
t = StepTracker(max_concurrent_batches=1)
t.add("h1", "req_a")
t.step(set()) # commit
result = t.step(set()) # expiry fires h1
assert result == ["h1"]
# req_a finishes later — must not re-fire.
result = t.step({"req_a"})
assert result == []
def test_multiplicity_same_hash_different_requests():
"""Two adds of the same mm_hash (different requests) → returned twice."""
t = StepTracker(max_concurrent_batches=3)
t.add("h1", "req_a")
t.add("h1", "req_b")
# Both requests finish in the same step.
result = t.step({"req_a", "req_b"})
assert sorted(result) == ["h1", "h1"]
def test_mixed_expiry_and_fast_path():
"""Some entries expire, others fast-path, in the same step."""
t = StepTracker(max_concurrent_batches=1)
t.add("h1", "req_a")
t.step(set()) # commit h1's slot
# Add h2 in the next step.
t.add("h2", "req_b")
# This step: h1's slot expires AND req_b finishes.
result = t.step({"req_b"})
assert set(result) == {"h1", "h2"}
def test_drain_all_returns_unprocessed():
"""drain_all returns everything pending and clears state."""
t = StepTracker(max_concurrent_batches=3)
t.add("h1", "req_a")
t.add("h2", "req_b")
t.step(set()) # commit both
t.add("h3", "req_c") # in _current, not yet committed
result = t.drain_all()
assert set(result) == {"h1", "h2", "h3"}
# After drain, everything is empty.
assert t.step(set()) == []
assert t.drain_all() == []
def test_drain_all_skips_already_processed():
"""drain_all should not re-return entries already fast-pathed."""
t = StepTracker(max_concurrent_batches=3)
t.add("h1", "req_a")
t.add("h2", "req_b")
# Fast-path h1.
result = t.step({"req_a"})
assert result == ["h1"]
# drain should only return h2.
result = t.drain_all()
assert result == ["h2"]
def test_empty_step_returns_nothing():
"""step() with nothing pending returns empty list."""
t = StepTracker(max_concurrent_batches=2)
assert t.step(set()) == []
assert t.step({"req_x"}) == []
def test_add_after_step_goes_to_next_slot():
"""Entries added after a step go into the next slot, not the current."""
t = StepTracker(max_concurrent_batches=1)
t.step(set()) # empty step
t.add("h1", "req_a")
# h1 is in _current. Next step commits it.
t.step(set())
# Now the deque is full (1 slot with h1). Next step expires it.
result = t.step(set())
assert result == ["h1"]
def test_fast_path_on_current_entries():
"""Entries still in _current (not yet committed) can be fast-pathed."""
t = StepTracker(max_concurrent_batches=3)
t.add("h1", "req_a")
# req_a finishes in the same step h1 was added.
result = t.step({"req_a"})
assert result == ["h1"]
# Entry is committed to slot but marked processed — won't re-fire.
for _ in range(5):
assert t.step(set()) == []
tests/v1/ec_connector/unit/cpu/worker/test_worker.py
+342 -71
View File
@@ -22,14 +22,16 @@ Mocking policy
import contextlib
import logging
import time
import uuid
from collections import deque
from unittest.mock import MagicMock, Mock, patch
import numpy as np
import pytest
import torch
from vllm.config import VllmConfig
from vllm.config.parallel import ParallelConfig
from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
)
@@ -55,6 +57,18 @@ _requires_accelerator = pytest.mark.skipif(
reason="exercises real accelerator stream/event coordination in ECCPUWorker",
)
# The memory-lifetime tests hold a copy back with `torch.cuda._sleep`, which the
# CUDA-like platforms provide but XPU does not.
_requires_cuda_alike = pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="stalling the compute stream requires a CUDA-like platform",
)
# Cycles to stall the compute stream for. Must comfortably outlast the host-side
# work the lifetime tests do while a copy is deliberately held back — ~0.5 s on
# a current datacenter GPU.
_STALL_CYCLES = 1_000_000_000
def _make_region() -> ECSharedRegion:
"""Fresh region backed by a real per-test mmap file."""
@@ -66,18 +80,28 @@ def _make_region() -> ECSharedRegion:
def _vllm_config(rank: int = 0) -> Mock:
cfg = Mock(spec=VllmConfig)
cfg.parallel_config = Mock(spec=ParallelConfig)
cfg.parallel_config.rank = rank
cfg.model_config = Mock()
cfg.model_config.dtype = _DTYPE
return cfg
return create_ec_vllm_config(rank=rank, dtype=_DTYPE)
def _meta(
*, saves: dict | None = None, loads: dict | None = None
) -> ECCPUConnectorMetadata:
return ECCPUConnectorMetadata(saves=saves or {}, loads=loads or {})
"""Build step metadata from plain `{mm_hash: block_ids}` dicts.
Load transfer ids are synthesized here so tests can keep naming loads by
mm_hash; `_load_id` recovers the id the worker will report.
"""
return ECCPUConnectorMetadata(
saves=saves or {},
loads={
mm_hash: (idx, block_ids)
for idx, (mm_hash, block_ids) in enumerate((loads or {}).items())
},
)
def _load_id(meta: ECCPUConnectorMetadata, mm_hash: str) -> int:
return meta.loads[mm_hash][0]
# ── fixtures ─────────────────────────────────────────────────────────────────
@@ -136,19 +160,52 @@ def make_worker():
region.cleanup()
def _wait_for_completion(
worker: ECCPUWorker, expected, direction: str, timeout_s: float = 30.0
) -> None:
"""Poll ``build_connector_worker_meta`` until ``expected`` is reported done.
Saves are reported by mm_hash, loads by transfer id.
"""
deadline = time.monotonic() + timeout_s
while time.monotonic() < deadline:
meta = worker.build_connector_worker_meta()
if meta is not None and expected in getattr(meta, f"completed_{direction}"):
return
time.sleep(0.001)
raise AssertionError(f"{direction} for {expected!r} never reported complete")
def _warm_up_stream_pools(worker: ECCPUWorker) -> None:
"""Run one throwaway save so the stream and event pools are populated.
The first CUDA stream creation in a process blocks the host for ~0.5 s on
driver initialization, which would swallow the stall the lifetime tests
below install. One flush puts a stream and its events in the worker's pools,
matching the steady state every step after the first sees.
"""
scratch = torch.zeros(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)
worker.save_caches({"w": scratch}, "w", _meta(saves={"w": [0]}))
worker.flush_saves()
_wait_for_completion(worker, "w", "saves")
assert worker._stream_pool, "warm-up did not recycle a stream"
# ── save_caches ──────────────────────────────────────────────────────────────
@_requires_accelerator
@pytest.mark.parametrize(
"n_elements,n_blocks",
"n_elements,block_ids",
[
(3 * _HIDDEN_DIM, 3), # exact fit: 3 full blocks, no remainder
(3 * _HIDDEN_DIM + 4, 4), # partial last block: 4 fp16 trail in slot 3
(3 * _HIDDEN_DIM, [7, 2, 5]), # exact fit: 3 full blocks, no remainder
(3 * _HIDDEN_DIM + 4, [7, 2, 5, 0]), # partial last block: 4 fp16 trail
(3 * _HIDDEN_DIM, [2, 3, 4]), # adjacent: one coalesced copy
(3 * _HIDDEN_DIM + 4, [2, 3, 4, 6]), # coalesced run, then a gap
],
ids=["exact-fit", "partial-last-block"],
ids=["exact-fit", "partial-last-block", "adjacent-blocks", "run-then-gap"],
)
def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks):
def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, block_ids):
"""``save_caches`` + ``flush_saves`` copies the source GPU tensor's bytes
into the block IDs named by ``meta.saves``, in the order they appear,
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
When ``total_bytes`` is not a multiple of ``block_size_bytes`` the last
block is partially written; the unwritten tail must remain whatever
was there before. Blocks not named in ``meta.saves`` must be untouched.
Blocks adjacent in the region share one descriptor, which must land the
same bytes in the same blocks as one descriptor per block does.
"""
worker = make_worker()
sentinel = 0x5A
@@ -165,9 +224,9 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks
expected_bytes = src.cpu().reshape(-1).view(torch.uint8)
total_bytes = n_elements * _DTYPE.itemsize
block_ids = [7, 2, 5, 0][:n_blocks]
worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))
worker.flush_saves()
_wait_for_completion(worker, "h", "saves")
for slot, block_idx in enumerate(block_ids):
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
)
@_requires_accelerator
@pytest.mark.parametrize(
"n_elements,block_ids,expected_sizes",
[
(3 * _HIDDEN_DIM, [2, 3, 4], [3 * _BLOCK_SIZE_BYTES]),
(3 * _HIDDEN_DIM, [7, 2, 5], [_BLOCK_SIZE_BYTES] * 3),
(3 * _HIDDEN_DIM, [2, 3, 6], [2 * _BLOCK_SIZE_BYTES, _BLOCK_SIZE_BYTES]),
(2 * _HIDDEN_DIM + 4, [2, 3, 4], [(2 * _HIDDEN_DIM + 4) * _DTYPE.itemsize]),
(_HIDDEN_DIM, [2, 3, 4], [_BLOCK_SIZE_BYTES]),
],
ids=[
"one-run",
"no-run",
"run-then-gap",
"run-stops-at-output",
"blocks-past-output-skipped",
],
)
def test_save_caches_coalesces_only_adjacent_blocks(
make_worker, n_elements, block_ids, expected_sizes
):
"""One descriptor per run of region-adjacent blocks, not one per block.
Per-descriptor overhead dominates these copies, so the byte tests above
would pass either way; this pins the descriptor layout that makes them
fast. A coalesced copy also stops at the encoder output, since entries are
sized by placeholder count, which can exceed the number of embeddings.
"""
worker = make_worker()
src = torch.arange(n_elements, dtype=_DTYPE, device=DEVICE_TYPE)
worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))
assert worker._save_bufs.sizes[: worker._save_count].tolist() == expected_sizes
def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker):
"""When the scheduler hasn't pre-allocated blocks for ``mm_hash``,
``save_caches`` + ``flush_saves`` is a pure no-op."""
@@ -245,6 +340,7 @@ def test_save_caches_batches_multiple_hashes(make_worker):
worker.save_caches(cache, "a", _meta(saves={"a": [1], "b": [3]}))
worker.save_caches(cache, "b", _meta(saves={"a": [1], "b": [3]}))
worker.flush_saves()
_wait_for_completion(worker, "a", "saves")
expected_a = src_a.cpu().view(torch.uint8)
expected_b = src_b.cpu().view(torch.uint8)
@@ -258,6 +354,50 @@ def test_save_caches_batches_multiple_hashes(make_worker):
assert torch.all(worker._region.blocks[idx] == sentinel)
class _HighAddress:
"""A device tensor reporting an address with the high bit set.
XPU USM allocations live up there and no host allocator will hand out such
an address, so naming one is the only way to reach the arithmetic. Every
attribute other than the address comes from the real tensor.
"""
def __init__(self, real: torch.Tensor, address: int):
self._real = real
self._address = address
def __getattr__(self, name):
return getattr(self._real, name)
def view(self, *args):
return self
def data_ptr(self) -> int:
return self._address
@_requires_accelerator
def test_save_descriptors_hold_high_bit_source_addresses(make_worker):
"""Source addresses >= 2**63 survive into the descriptor buffers."""
worker = make_worker()
real = torch.zeros(2 * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)
address = 2**64 - 8 * _BLOCK_SIZE_BYTES
# Two runs, so the second descriptor's address is an offset off the base.
worker.save_caches(
{"a": _HighAddress(real, address)}, "a", _meta(saves={"a": [1, 5]})
)
assert worker._save_count == 2
stored = worker._save_bufs.src_np[:2].astype(np.uint64)
assert list(stored) == [address, address + _BLOCK_SIZE_BYTES]
# The batch describes an address nothing may read from; drop it unflushed.
worker._save_bufs = None
worker._save_count = 0
worker._save_mm_hashes = []
# ── start_load_caches ────────────────────────────────────────────────────────
@@ -288,19 +428,31 @@ def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker
@_requires_accelerator
def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker):
"""If ``encoder_cache`` already holds the ``mm_hash``, the worker must
not overwrite it."""
@pytest.mark.parametrize(
"loads,expected_sizes",
[
({"a": [2, 3], "b": [4, 5]}, [4 * _BLOCK_SIZE_BYTES]),
({"b": [4, 5], "a": [2, 3]}, [4 * _BLOCK_SIZE_BYTES]),
({"a": [2, 3], "b": [6, 7]}, [2 * _BLOCK_SIZE_BYTES] * 2),
],
ids=["adjacent-entries", "adjacent-entries-reversed", "gap-between-entries"],
)
def test_start_load_caches_coalesces_across_entries(make_worker, loads, expected_sizes):
"""Blocks adjacent in the region share a descriptor across entries too.
Entries are ordered by first block, so whether they coalesce follows from
the region layout rather than the order the scheduler dispatched them.
"""
worker = make_worker()
worker._region.blocks[0].fill_(0x42)
captured: list[list[int]] = []
sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device=DEVICE_TYPE)
encoder_cache = {"h": sentinel}
worker.start_load_caches(encoder_cache, _meta(loads={"h": [0]}))
with patch(
"vllm.distributed.ec_transfer.ec_connector.cpu.worker.swap_blocks_batch",
side_effect=lambda src, dst, sizes, **kw: captured.append(sizes.tolist()),
):
worker.start_load_caches({}, _meta(loads=loads))
assert encoder_cache["h"] is sentinel, (
"existing encoder_cache entry must not be replaced"
)
assert captured == [expected_sizes]
@_requires_accelerator
@@ -315,30 +467,34 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker):
@_requires_accelerator
def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):
"""Cached entries are preserved while new ones are loaded."""
"""Every hash in ``meta.loads`` is copied from mmap, including one whose
key is already resident in ``encoder_cache`` — there is no worker-side
skip."""
worker = make_worker()
n_blocks = 2
n_blocks = 3
src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape(
n_blocks, _HIDDEN_DIM
)
src_int8 = src_orig.view(torch.int8).reshape(n_blocks, _BLOCK_SIZE_BYTES)
new_block_ids = [4, 2]
for i, idx in enumerate(new_block_ids):
block_ids = [0, 4, 2]
for i, idx in enumerate(block_ids):
worker._region.blocks[idx].copy_(src_int8[i])
cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device=DEVICE_TYPE)
encoder_cache: dict[str, torch.Tensor] = {"cached_h": cached_tensor}
worker.start_load_caches(
encoder_cache,
_meta(loads={"cached_h": [0], "new_h": new_block_ids}),
_meta(loads={"resident_h": [block_ids[0]], "new_h": block_ids[1:]}),
)
assert encoder_cache["cached_h"] is cached_tensor
assert "new_h" in encoder_cache
out = encoder_cache["new_h"]
assert out.shape == (n_blocks, _HIDDEN_DIM)
assert out.dtype == _DTYPE
assert torch.equal(out.cpu(), src_orig)
resident = encoder_cache["resident_h"]
assert resident is not cached_tensor, "resident entry must be reloaded, not skipped"
assert torch.equal(resident.cpu(), src_orig[:1])
new = encoder_cache["new_h"]
assert new.shape == (2, _HIDDEN_DIM)
assert new.dtype == _DTYPE
assert torch.equal(new.cpu(), src_orig[1:])
@_requires_accelerator
@@ -382,6 +538,7 @@ def test_save_then_load_round_trips_bytes(make_worker):
).reshape(n_blocks, _HIDDEN_DIM)
worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))
worker.flush_saves()
_wait_for_completion(worker, "h", "saves")
encoder_cache: dict[str, torch.Tensor] = {}
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):
assert torch.equal(out.cpu(), src.cpu())
# ── memory lifetime across in-flight copies ─────────────────────────────────
@_requires_cuda_alike
def test_save_survives_encoder_cache_free_before_copy_runs(make_worker):
"""The bytes handed to ``save_caches`` must reach the mmap intact even when
the caller drops its ``encoder_cache`` reference before the copy has run.
The model runner pops encoder cache entries on whatever step the scheduler
asks it to, with no dependency on the save copy having landed. Dropping the
last reference hands the GPU memory back to the caching allocator, which is
free to reuse it for the next same-sized allocation on the compute stream —
the copy reads that memory on a different stream, so the allocator has no
reason to hold it back.
The stall makes the hazard deterministic rather than timing-dependent: the
save copy is gated behind the compute stream, so it provably has not started
while the overwrite (issued on an ungated stream and waited on) completes.
The mmap must be pinned, as it is in production: a copy into pageable host
memory blocks the host until it lands, which would hide the hazard.
"""
worker = make_worker(pin_memory_available=True)
_warm_up_stream_pools(worker)
worker._region.blocks.fill_(0x5A)
block_ids = [0, 1, 2]
saved = torch.full(
(len(block_ids), _HIDDEN_DIM), 1.0, dtype=_DTYPE, device=DEVICE_TYPE
)
expected = saved.cpu().reshape(-1).view(torch.uint8)
encoder_cache = {"h": saved}
worker.save_caches(encoder_cache, "h", _meta(saves={"h": block_ids}))
torch.cuda._sleep(_STALL_CYCLES)
worker.flush_saves()
saved_ptr = saved.data_ptr()
del saved
encoder_cache.clear()
# Same size on the same stream, so the allocator hands back the block just
# freed. Overwriting it from an ungated stream — and waiting for that write
# — orders the overwrite strictly before the still-stalled save copy.
overwrite = torch.empty(
(len(block_ids), _HIDDEN_DIM), dtype=_DTYPE, device=DEVICE_TYPE
)
reused = overwrite.data_ptr() == saved_ptr
other_stream = torch.cuda.Stream()
with torch.cuda.stream(other_stream):
overwrite.fill_(-1.0)
other_stream.synchronize()
_wait_for_completion(worker, "h", "saves")
actual = worker._region.blocks[block_ids].reshape(-1).view(torch.uint8)
assert torch.equal(actual, expected), (
"mmap holds bytes written after save_caches was called; the copy read "
f"freed source memory (allocator reused the freed block: {reused})"
)
@_requires_cuda_alike
def test_load_buffer_survives_eviction_while_consumer_read_is_queued(make_worker):
"""A loaded encoder cache entry must keep its bytes for a consumer already
queued on the compute stream, even after the entry is dropped and another
load is dispatched.
``start_load_caches`` allocates its destination inside the load stream's
context, which ties that memory to the load stream in the caching allocator.
Dropping the encoder cache entry returns it to the load stream's pool, where
the next load can claim it — while the model's read of it is still queued on
the compute stream.
"""
worker = make_worker(pin_memory_available=True)
_warm_up_stream_pools(worker)
a_blocks, b_blocks = [4, 5], [6, 7]
for block_idx, fill in zip(a_blocks + b_blocks, (0x11, 0x12, 0x21, 0x22)):
worker._region.blocks[block_idx].fill_(fill)
expected_a = worker._region.blocks[a_blocks].reshape(-1).clone()
encoder_cache: dict[str, torch.Tensor] = {}
meta_a = _meta(loads={"a": a_blocks})
worker.start_load_caches(encoder_cache, meta_a)
loaded_a = encoder_cache["a"]
# Queue a consumer of the loaded entry on the compute stream behind a stall,
# the way the model reads the encoder cache later in the step.
torch.cuda._sleep(_STALL_CYCLES)
consumed = loaded_a.clone()
# Draining the completion report recycles the load stream into the pool, so
# the next load allocates from the same allocator pool the entry was freed to.
_wait_for_completion(worker, _load_id(meta_a, "a"), "loads")
loaded_a_ptr = loaded_a.data_ptr()
del loaded_a
encoder_cache.clear()
worker.start_load_caches(encoder_cache, _meta(loads={"b": b_blocks}))
reused = encoder_cache["b"].data_ptr() == loaded_a_ptr
torch.accelerator.synchronize()
assert torch.equal(consumed.cpu().reshape(-1).view(torch.int8), expected_a), (
"queued consumer read bytes from a later load; the load destination was "
f"recycled while still in use (allocator reused the buffer: {reused})"
)
# ── buffer recycling ────────────────────────────────────────────────────────
@_requires_accelerator
def test_buffer_pool_is_reused_across_save_steps(make_worker):
"""After flush_saves, descriptor buffers are returned to the pool and
reused on the next flush — no reallocation."""
"""Once a save copy completes its descriptor buffers return to the pool and
are reused by the next flush — no reallocation.
The buffers hold the addresses the copy reads, so they stay out of the pool
until the transfer's end event fires.
"""
worker = make_worker()
src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)
worker.save_caches({"h": src}, "h", _meta(saves={"h": [0]}))
worker.flush_saves()
_wait_for_completion(worker, "h", "saves")
assert len(worker._buf_pool._pool) == 1
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):
# Second step reuses the same buffer.
worker.save_caches({"h": src}, "h", _meta(saves={"h": [1]}))
worker.flush_saves()
_wait_for_completion(worker, "h", "saves")
assert len(worker._buf_pool._pool) == 1
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):
@_requires_accelerator
def test_buffer_pool_is_reused_across_load_steps(make_worker):
"""After start_load_caches, descriptor buffers are returned to the pool
and reused on the next call."""
"""Once a load copy completes its descriptor buffers return to the pool and
are reused by the next call — no reallocation."""
worker = make_worker()
worker._region.blocks[0].fill_(0x01)
worker._region.blocks[1].fill_(0x02)
encoder_cache: dict[str, torch.Tensor] = {}
worker.start_load_caches(encoder_cache, _meta(loads={"a": [0]}))
meta = _meta(loads={"a": [0]})
worker.start_load_caches(encoder_cache, meta)
_wait_for_completion(worker, _load_id(meta, "a"), "loads")
assert len(worker._buf_pool._pool) == 1
buf_id = id(worker._buf_pool._pool[0].src_ptrs)
encoder_cache2: dict[str, torch.Tensor] = {}
worker.start_load_caches(encoder_cache2, _meta(loads={"b": [1]}))
meta2 = _meta(loads={"b": [1]})
worker.start_load_caches(encoder_cache2, meta2)
_wait_for_completion(worker, _load_id(meta2, "b"), "loads")
assert len(worker._buf_pool._pool) == 1
assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id
# ── stream management ────────────────────────────────────────────────────────
@_requires_accelerator
def test_stream_initialized_at_construction(make_worker):
"""``_load_stream`` must be a fully initialized accelerator stream."""
worker = make_worker()
assert isinstance(worker._load_stream, current_platform.Stream)
# ── lifecycle ────────────────────────────────────────────────────────────────
@@ -500,12 +767,15 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):
worker = object.__new__(ECCPUWorker)
mock_region = Mock(spec=ECSharedRegion)
worker._region = mock_region
worker._load_stream = MagicMock()
worker._save_bufs = None
worker._save_count = 0
worker._save_mm_hashes = []
worker._inflight_saves = deque()
worker._inflight_loads = deque()
worker._stream_pool = []
worker._event_pool = []
worker.shutdown()
worker._load_stream.synchronize.assert_called_once()
mock_region.cleanup.assert_called_once()
mock_region.cleanup.side_effect = RuntimeError("boom")
@@ -525,12 +795,13 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):
@_requires_accelerator
def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):
"""Full pipeline: scheduler allocates blocks, worker saves GPU tensor to
mmap via flush_saves, scheduler marks ready after step delay, worker
loads from mmap back to GPU, and the result matches the original.
"""Full pipeline: scheduler allocates blocks, worker saves a GPU tensor to
mmap via flush_saves, the worker's completion report marks the entry ready,
the worker loads from mmap back to GPU, and the result matches the original.
Exercises the real scheduler + worker cooperation through a shared
ECSharedRegion, with real accelerator transfers and stream coordination.
ECSharedRegion, with real accelerator transfers and stream coordination,
and the event-driven mark_ready path.
"""
import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
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):
# Build scheduler sharing the same region.
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)
class _EC:
is_ec_producer = True
is_ec_consumer = True
class _Cfg:
ec_transfer_config = _EC()
max_concurrent_batches = 1
scheduler = ECCPUScheduler(_Cfg())
scheduler = ECCPUScheduler(_vllm_config())
# -- Step 1: scheduler allocates, worker saves --
n_blocks = 3
@@ -580,16 +843,24 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):
worker.save_caches(encoder_cache, "img_001", meta_save)
worker.flush_saves()
# -- Step 2: readiness delay (max_concurrent_batches=1) --
meta_step2 = scheduler.build_connector_meta(scheduler_output=None)
# -- Step 2: worker reports the save memcpy complete → scheduler marks ready --
torch.accelerator.synchronize()
worker_meta = worker.build_connector_worker_meta()
assert worker_meta is not None
assert "img_001" in worker_meta.completed_saves
class _Output:
ec_connector_worker_meta = worker_meta
scheduler.update_connector_output(_Output())
assert scheduler.has_cache_item("img_001") is True
assert meta_step2.loads == {} # no load requested yet
# -- Step 3: scheduler emits load, worker loads --
scheduler.update_state_after_alloc(_Request(), 0)
meta_load = scheduler.build_connector_meta(scheduler_output=None)
assert "img_001" in meta_load.loads
assert meta_load.loads["img_001"] == meta_save.saves["img_001"]
_, load_blocks = meta_load.loads["img_001"]
assert load_blocks == meta_save.saves["img_001"]
load_cache: dict[str, torch.Tensor] = {}
worker.start_load_caches(load_cache, meta_load)
tests/v1/ec_connector/unit/test_ec_cpu_connector.py
+4 -3
View File
@@ -48,9 +48,10 @@ def _make_unique_image(base_image, iteration: int):
def _wait_for_ec_ready(llm: LLM) -> None:
"""Send a dummy request to force a scheduler step.
After generate() returns, the EC entry's data is in the mmap but
readiness may not have fired yet (first-finish fires in the NEXT
scheduler step). This forces that step.
When generate() returns, the save copy has been issued to the GPU but its
completion event may fire on a later step, so the entry is not yet marked
ready. This extra step lets the scheduler poll the worker's completion
report and mark the entry ready.
"""
llm.generate("hi", SamplingParams(max_tokens=1), use_tqdm=False)
tests/v1/ec_connector/unit/test_metadata.py
+13 -13
View File
@@ -38,19 +38,19 @@ def test_metadata_saves_dict_operations():
assert meta.saves["mm_hash_2"] == [10, 11]
def test_metadata_loads_dict_operations():
"""The loads dict supports standard dict operations."""
def test_metadata_loads_carry_a_transfer_id_and_blocks():
"""Each load entry is a (transfer_id, block_ids) pair, so the worker can
report the transfer rather than the mm_hash."""
meta = ECCPUConnectorMetadata()
# Add multiple entries
meta.loads["mm_hash_1"] = [5, 6, 7]
meta.loads["mm_hash_2"] = [100]
meta.loads["mm_hash_1"] = (0, [5, 6, 7])
meta.loads["mm_hash_2"] = (1, [100])
assert "mm_hash_1" in meta.loads
assert "mm_hash_2" in meta.loads
assert len(meta.loads) == 2
assert meta.loads["mm_hash_1"] == [5, 6, 7]
assert meta.loads["mm_hash_2"] == [100]
assert meta.loads["mm_hash_1"] == (0, [5, 6, 7])
transfer_id, block_ids = meta.loads["mm_hash_2"]
assert transfer_id == 1
assert block_ids == [100]
def test_metadata_saves_and_loads_are_independent():
@@ -58,11 +58,11 @@ def test_metadata_saves_and_loads_are_independent():
meta = ECCPUConnectorMetadata()
meta.saves["key"] = [1, 2]
meta.loads["key"] = [3, 4]
meta.loads["key"] = (0, [3, 4])
# Same key but different dicts should not interfere
assert meta.saves["key"] == [1, 2]
assert meta.loads["key"] == [3, 4]
assert meta.loads["key"] == (0, [3, 4])
# Modifying one should not affect the other
meta.saves["other"] = [99]
@@ -74,11 +74,11 @@ def test_metadata_block_indices_are_integers():
meta = ECCPUConnectorMetadata()
meta.saves["hash"] = [0, 100, 999, 65535]
meta.loads["hash"] = [1, 2, 3]
meta.loads["hash"] = (0, [1, 2, 3])
for idx in meta.saves["hash"]:
assert isinstance(idx, int)
for idx in meta.loads["hash"]:
for idx in meta.loads["hash"][1]:
assert isinstance(idx, int)
tests/v1/ec_connector/unit/utils.py
+70
View File
@@ -0,0 +1,70 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared config builders for EC connector unit tests."""
import uuid
from unittest.mock import Mock
import torch
from vllm.config import ModelConfig, VllmConfig
from vllm.config.ec_transfer import ECRole, ECTransferConfig
from vllm.config.parallel import ParallelConfig
# The scheduler derives its per-load ack count from the PCP size only under the
# `mp` backend -- which also keeps the world size off the local GPU count, so
# TP > 1 stays expressible on a single-GPU host.
_EXECUTOR_BACKEND = "mp"
_PCP_SIZE = 1
def create_ec_vllm_config(
*,
ec_role: ECRole = "ec_both",
tensor_parallel_size: int = 1,
rank: int = 0,
dtype: torch.dtype = torch.float16,
) -> Mock:
"""Build a `VllmConfig` stand-in for EC connector unit tests.
`ec_transfer_config` and `parallel_config` are the real config objects, so
role derivation and the rank arithmetic the connector depends on behave as
they do in production.
`model_config` is a stub because nothing under test needs a real one: the
connector reads only `dtype` here, and the paths that inspect the model
itself are not exercised. Building a real `ModelConfig` would resolve an HF
model and inspect its architecture.
Args:
ec_role: EC role this instance plays. `is_ec_producer` and
`is_ec_consumer` are derived from it, as in production.
tensor_parallel_size: TP degree.
rank: Global rank of this worker.
dtype: Encoder cache dtype.
Returns:
A `VllmConfig`-specced mock carrying the configs above.
"""
parallel_config = ParallelConfig(
tensor_parallel_size=tensor_parallel_size,
prefill_context_parallel_size=_PCP_SIZE,
distributed_executor_backend=_EXECUTOR_BACKEND,
rank=rank,
)
ec_transfer_config = ECTransferConfig(
# The role properties report False unless a connector is named.
ec_connector="ECCPUConnector",
ec_role=ec_role,
engine_id=str(uuid.uuid4()),
)
model_config = Mock(spec=ModelConfig)
model_config.dtype = dtype
vllm_config = Mock(spec=VllmConfig)
vllm_config.ec_transfer_config = ec_transfer_config
vllm_config.parallel_config = parallel_config
vllm_config.model_config = model_config
vllm_config.instance_id = f"ec-test-{uuid.uuid4()}"
return vllm_config
vllm/distributed/ec_transfer/ec_connector/cpu/common.py
+36 -5
View File
@@ -7,7 +7,10 @@ from typing import TYPE_CHECKING
import torch
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
from vllm.distributed.ec_transfer.ec_connector.base import (
ECConnectorMetadata,
ECConnectorWorkerMetadata,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
)
@@ -29,10 +32,38 @@ class ECCPUConnectorMetadata(ECConnectorMetadata):
# encoder_cache[mm_hash] → mmap at these block IDs.
saves: dict[str, list[int]] = field(default_factory=dict)
# Consumer role: mm_hashes whose bytes are available in the local mmap;
# the worker's start_load_caches copies mmap[block_ids] → GPU
# encoder_cache.
loads: dict[str, list[int]] = field(default_factory=dict)
# Consumer role: mm_hashes whose bytes are available in the local mmap,
# mapped to (transfer_id, block_ids); the worker's start_load_caches
# copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id
# once its copy lands. The id is minted per dispatch, so two loads of the
# same mm_hash in different steps are distinct and a late report from the
# earlier one cannot release the later one's pin.
loads: dict[str, tuple[int, list[int]]] = field(default_factory=dict)
@dataclass
class ECCPUWorkerMetadata(ECConnectorWorkerMetadata):
"""Per-step worker → scheduler payload for the ECCPUConnector.
Reports the GPU copies that completed this step: saved mm_hashes become
safe to mark ready, and loaded transfers become safe to unpin once every
participating rank has reported them. Built by
`ECCPUWorker.build_connector_worker_meta`; consumed by
`ECCPUScheduler.update_connector_output`.
Loads are reported by transfer id rather than mm_hash because every rank
copies the same blocks: `aggregate` concatenates, so an id appears once
per reporting rank and the scheduler can count participants off the list.
"""
completed_saves: list[str] = field(default_factory=list)
completed_loads: list[int] = field(default_factory=list)
def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetadata:
assert isinstance(other, ECCPUWorkerMetadata)
self.completed_saves.extend(other.completed_saves)
self.completed_loads.extend(other.completed_loads)
return self
def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int:
vllm/distributed/ec_transfer/ec_connector/cpu/connector.py
+16
View File
@@ -15,6 +15,7 @@ import torch
from vllm.distributed.ec_transfer.ec_connector.base import (
ECConnectorBase,
ECConnectorRole,
ECConnectorWorkerMetadata,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
@@ -24,6 +25,7 @@ from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ECConnectorOutput
from vllm.v1.request import Request
logger = init_logger(__name__)
@@ -113,6 +115,20 @@ class ECCPUConnector(ECConnectorBase):
self.connector_worker.flush_saves()
return None, None
def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None:
if self.connector_worker is not None:
return self.connector_worker.build_connector_worker_meta()
return None
def update_connector_output(self, connector_output: "ECConnectorOutput") -> None:
if self.connector_scheduler is not None:
self.connector_scheduler.update_connector_output(connector_output)
def has_pending_push_work(self) -> bool:
if self.connector_scheduler is not None:
return self.connector_scheduler.has_pending_push_work()
return False
# Shared.
def shutdown(self) -> None:
if self.connector_scheduler is not None:
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
+64 -32
View File
@@ -11,19 +11,18 @@ from typing import TYPE_CHECKING
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
ECCPUWorkerMetadata,
create_ec_shared_region,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
EmbeddingCache,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.step_tracker import (
StepTracker,
)
from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ECConnectorOutput
from vllm.v1.request import Request
logger = init_logger(__name__)
@@ -42,16 +41,27 @@ class ECCPUScheduler:
# Block allocator + LRU eviction policy for the shared region.
self._cache = EmbeddingCache(self._region.num_blocks)
max_batches = vllm_config.max_concurrent_batches
# Delays mark_ready until the GPU→mmap DMA is guaranteed complete.
self._ready_tracker = StepTracker(max_batches)
# Delays unpin until the mmap→GPU DMA is guaranteed complete.
self._unpin_tracker = StepTracker(max_batches)
# mm_hash → block IDs allocated this step for GPU→mmap saves.
self._pending_saves: dict[str, list[int]] = {}
# mm_hash → block IDs to load from mmap→GPU this step.
self._pending_loads: dict[str, list[int]] = {}
# mm_hash → (transfer_id, block IDs) to load from mmap→GPU this step.
self._pending_loads: dict[str, tuple[int, list[int]]] = {}
# Dispatched loads awaiting completion reports, keyed by transfer id:
# transfer_id → (mm_hash, reports still outstanding). The pin taken at
# dispatch is released only once the count reaches zero.
self._load_acks: dict[int, tuple[str, int]] = {}
self._next_transfer_id = 0
pc = vllm_config.parallel_config
# Reports to expect per load. Only pipeline stage 0 runs the EC
# connector, so the participants are the tp × pcp ranks of that stage.
# Other executors deliver a single rank's output, and expecting reports
# that cannot arrive would hold the pin forever.
self._expected_load_reports = (
pc.tensor_parallel_size * pc.prefill_context_parallel_size
if pc.distributed_executor_backend == "mp"
else 1
)
def has_cache_item(self, identifier: str) -> bool:
if not self._is_consumer:
@@ -72,30 +82,19 @@ class ECCPUScheduler:
entry = self._cache.alloc(mm_hash, feature.mm_position.length)
if entry is not None:
self._pending_saves[mm_hash] = list(entry.block_ids)
self._ready_tracker.add(mm_hash, request.request_id)
if self._is_consumer and mm_hash not in self._pending_loads:
entry = self._cache.get(mm_hash)
if entry is not None and entry.ready:
self._cache.pin(mm_hash)
self._pending_loads[mm_hash] = list(entry.block_ids)
self._unpin_tracker.add(mm_hash, request.request_id)
transfer_id = self._next_transfer_id
self._next_transfer_id += 1
self._pending_loads[mm_hash] = (transfer_id, list(entry.block_ids))
self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports)
def build_connector_meta(
self, scheduler_output: "SchedulerOutput"
) -> ECCPUConnectorMetadata:
finished = (
scheduler_output.finished_req_ids if scheduler_output is not None else set()
)
for key in self._ready_tracker.step(finished):
entry = self._cache.get(key)
if entry is not None and not entry.ready:
self._cache.mark_ready(key)
for key in self._unpin_tracker.step(finished):
self._cache.unpin(key)
meta = ECCPUConnectorMetadata()
if self._is_producer:
meta.saves = self._pending_saves
@@ -105,13 +104,46 @@ class ECCPUScheduler:
self._pending_loads = {}
return meta
def shutdown(self) -> None:
# drain_all() covers both entries still in _current (never
# consumed by build_connector_meta) and entries in slots.
self._pending_loads.clear()
for mm_hash in self._unpin_tracker.drain_all():
def update_connector_output(self, connector_output: "ECConnectorOutput") -> None:
"""Apply the worker's memcpy-completion report to the cache.
Completed saves become safe to mark ready. A load is unpinned once
every participating rank has reported its transfer id; reports for a
transfer that has already been released, or for one this scheduler
never dispatched, are ignored.
"""
meta = connector_output.ec_connector_worker_meta
if not isinstance(meta, ECCPUWorkerMetadata):
return
for mm_hash in meta.completed_saves:
entry = self._cache.get(mm_hash)
if entry is not None and not entry.ready:
self._cache.mark_ready(mm_hash)
for transfer_id in meta.completed_loads:
pending = self._load_acks.get(transfer_id)
if pending is None:
continue
mm_hash, outstanding = pending
if outstanding > 1:
self._load_acks[transfer_id] = (mm_hash, outstanding - 1)
continue
# Drop the entry before unpinning so a replayed report is treated
# as stale rather than releasing the pin a second time.
del self._load_acks[transfer_id]
self._cache.unpin(mm_hash)
self._ready_tracker.drain_all()
def has_pending_push_work(self) -> bool:
"""True while any dispatched save or load has not been confirmed done.
Keeps the engine stepping so the worker's completion reports are polled
even when no requests are otherwise runnable.
"""
return self._cache.has_held_entries()
def shutdown(self) -> None:
self._pending_saves.clear()
self._pending_loads.clear()
self._load_acks.clear()
self._is_producer = False
self._is_consumer = False
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py
+13 -1
View File
@@ -67,6 +67,15 @@ class EmbeddingCache:
with self._lock:
return self._entries.get(key)
def has_held_entries(self) -> bool:
"""True if any entry is not-ready or pinned (an in-flight transfer).
Held entries are exactly those absent from the eviction free list, so
they represent saves awaiting mark_ready or loads awaiting unpin.
"""
with self._lock:
return len(self._entries) != len(self._entries_free_list)
def alloc(self, key: str, n_blocks: int) -> CacheEntry | None:
"""Allocate *n_blocks* for *key*, evicting as needed.
@@ -84,7 +93,10 @@ class EmbeddingCache:
if len(self._free_blocks) + self._evictable_block_count < n_blocks:
return None
self._evict_until(n_blocks)
block_ids = tuple(self._free_blocks.pop() for _ in range(n_blocks))
# Canonical ascending order: keeps a freed entry's blocks contiguous
# in the free list when it is evicted, and lets consumers detect
# which of an entry's blocks are adjacent.
block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks)))
entry = CacheEntry(block_ids)
self._entries[key] = entry
return entry
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/step_tracker.py
-115
View File
@@ -1,115 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""StepTracker — step-counting delay with first-finish fast-path.
Manages a collection of (mm_hash, request_id) entries that need to be
processed after a configurable number of engine steps have elapsed.
Additionally, if the originating request finishes before the step-count
expires, the entry is processed immediately (the GPU transfer is
guaranteed complete by request finish time).
Used by ECCPUScheduler: one instance for readiness marking, one for
unpins.
"""
from collections import deque
from dataclasses import dataclass
@dataclass(slots=True)
class _PendingEntry:
mm_hash: str
request_id: str
processed: bool = False
class StepTracker:
"""Track pending operations with step-count delay and first-finish.
Each ``add(mm_hash, request_id)`` registers an operation that will be
returned from ``step()`` either when:
(a) ``max_concurrent_batches`` steps elapse (deque expiry), OR
(b) ``request_id`` appears in ``finished_req_ids`` (first-finish).
Each add is processed exactly once — never both (a) and (b).
"""
def __init__(self, max_concurrent_batches: int) -> None:
# One slot per past step; oldest at the right. When the deque is
# full the rightmost slot is popped and its entries expire.
self._slots: deque[list[_PendingEntry]] = deque(maxlen=max_concurrent_batches)
# Reverse index: request_id → entries across all slots + _current.
# Enables O(1) lookup for the first-finish fast-path.
self._req_index: dict[str, list[_PendingEntry]] = {}
# Entries added this step (via add()), committed to a new slot at
# the start of the next step() call.
self._current: list[_PendingEntry] = []
def add(self, mm_hash: str, request_id: str) -> None:
"""Register an mm_hash needing processing after a delay."""
entry = _PendingEntry(mm_hash=mm_hash, request_id=request_id)
self._current.append(entry)
self._req_index.setdefault(request_id, []).append(entry)
def step(self, finished_req_ids: set[str]) -> list[str]:
"""Advance one step. Return mm_hashes ready for processing.
Returns:
List of mm_hashes to process. May contain duplicates if the
same mm_hash was added multiple times (from different requests).
"""
result: list[str] = []
# Phase 1: first-finish fast-path.
for req_id in finished_req_ids:
entries = self._req_index.pop(req_id, None)
if entries is None:
continue
for entry in entries:
if not entry.processed:
entry.processed = True
result.append(entry.mm_hash)
# Phase 2: step-count expiry.
if len(self._slots) == self._slots.maxlen:
expired_slot = self._slots.pop()
for entry in expired_slot:
if not entry.processed:
entry.processed = True
result.append(entry.mm_hash)
self._cleanup_expired(expired_slot)
# Phase 3: commit current entries to new slot.
self._slots.appendleft(self._current)
self._current = []
return result
def drain_all(self) -> list[str]:
"""Return all unprocessed entries and clear internal state."""
result: list[str] = []
for entry in self._current:
if not entry.processed:
entry.processed = True
result.append(entry.mm_hash)
self._current = []
for slot in self._slots:
for entry in slot:
if not entry.processed:
entry.processed = True
result.append(entry.mm_hash)
self._slots.clear()
self._req_index.clear()
return result
def _cleanup_expired(self, slot: list[_PendingEntry]) -> None:
"""Prune reverse index for entries in an expired slot."""
for entry in slot:
entries_for_req = self._req_index.get(entry.request_id)
if entries_for_req is None:
# Already removed by fast-path (pop in step()).
continue
entries_for_req.remove(entry)
if not entries_for_req:
del self._req_index[entry.request_id]
vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py
+266 -50
View File
@@ -7,13 +7,19 @@ per-step connector metadata (`ECCPUConnectorMetadata`) to decide which
blocks to copy in each direction.
"""
from collections import deque
from collections.abc import Sequence
from dataclasses import dataclass
from itertools import chain
from typing import TYPE_CHECKING
import numpy as np
import torch
from vllm._custom_ops import swap_blocks_batch
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
ECCPUWorkerMetadata,
create_ec_shared_region,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.worker.descriptor_buffers import (
@@ -34,6 +40,50 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
def _coalesce_runs(
block_ids: Sequence[int],
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Split block ids into runs that are consecutive in the region.
Returns `(slots, first_blocks, num_blocks)`, one element per run, where
`slots` is the run's offset in the flat block sequence. Per-descriptor
overhead dominates these copies, so describing a run with one descriptor
rather than one per block is worth an order of magnitude.
Both callers derive the non-region side of a run from its slot, so that
side advances in lockstep with the region and only the region has to be
split. `EmbeddingCache` hands out ascending ids, which is what makes runs
findable at all.
"""
ids = np.fromiter(block_ids, dtype=np.int64, count=len(block_ids))
breaks = np.flatnonzero(np.diff(ids) != 1) + 1
slots = np.concatenate(([0], breaks))
ends = np.concatenate((breaks, [ids.size]))
return slots, ids[slots], ends - slots
@dataclass
class Transfer:
"""A batched copy in flight on a dedicated GPU stream.
Its descriptor buffers and stream are held until `end_event` fires; only
then are the buffers and stream safe to recycle and the `completions` safe
to report. `start_event`/`end_event` bracket the copy so its elapsed time
can be measured on completion.
`completions` is what the scheduler is told about this copy: mm_hashes for
saves, transfer ids for loads.
"""
start_event: torch.Event
end_event: torch.Event
completions: list[str] | list[int]
bufs: DescriptorBuffers
stream: torch.Stream
num_bytes: int
class ECCPUWorker:
"""Worker-side delegate for the ECCPUConnector.
@@ -43,8 +93,18 @@ class ECCPUWorker:
batched call in `flush_saves`.
- Consumer role: copies `mmap[block_ids]` → `encoder_cache[mm_hash]`
for all entries in `metadata.loads` via a single `swap_blocks_batch`
call on the load stream.
call on a dedicated stream.
- On `ec_both` nodes both paths run back-to-back in a single step.
Every batched copy runs on a stream drawn from a pool so it overlaps the
compute stream, bracketed by a start/end event pair. The stream, events,
and descriptor buffers are recycled once the end event fires.
Both paths move memory between the two streams, so each buffer is registered
against the stream that is not its own: save sources against the save stream,
load destinations against the compute stream. Without that, the caching
allocator would be free to reissue memory a copy or a queued read still
needs, since it only tracks the stream a buffer was allocated on.
"""
def __init__(self, vllm_config: "VllmConfig") -> None:
@@ -62,15 +122,72 @@ class ECCPUWorker:
get_tensor_model_parallel_rank() == 0 and get_pcp_group().rank_in_group == 0
)
# Dedicated stream for async mmap→GPU loads (overlaps with compute).
self._load_stream = current_platform.Stream()
# Descriptor buffer pool (recycled across steps, shared by both paths).
self._buf_pool = DescriptorBufferPool()
# Active save buffer being filled during save_caches calls this step.
self._save_bufs: DescriptorBuffers | None = None
self._save_count: int = 0
# Bytes accumulated into the active save buffer this step, for logging.
self._save_bytes: int = 0
# mm_hashes accumulated into the active save buffer this step.
self._save_mm_hashes: list[str] = []
# Stream the active save buffer's copy will run on, held from the moment
# the buffer opens so `save_caches` can register its sources against it.
self._save_stream: torch.Stream | None = None
# Batched copies whose end event has not yet fired. Buffers and stream
# stay held until the event fires; the completions are reported to the
# scheduler only then. Separate per direction because saves and loads
# can be in flight concurrently.
self._inflight_saves: deque[Transfer] = deque()
self._inflight_loads: deque[Transfer] = deque()
# Streams and timing events recycled across steps to avoid per-step
# create/destroy churn. Shared by both directions.
self._stream_pool: list[torch.Stream] = []
self._event_pool: list[torch.Event] = []
def _acquire_stream(self) -> torch.Stream:
if self._stream_pool:
return self._stream_pool.pop()
return current_platform.Stream()
def _acquire_event(self) -> torch.Event:
return (
self._event_pool.pop()
if self._event_pool
else torch.Event(enable_timing=True)
)
def _collect_finished(self, inflight: deque[Transfer], direction: str) -> list:
"""Pop transfers whose end event has fired, recycle their stream,
events, and buffers, and return their completions.
The front check is conservative: a transfer is reported only once its
own end event fires, so completions are always genuinely done, and each
is reported exactly once because the transfer is popped as it is
reported. A later transfer that happens to finish first simply waits
behind the front and is reported on a subsequent poll.
"""
done: list = []
while inflight and inflight[0].end_event.query():
transfer = inflight.popleft()
done.extend(transfer.completions)
elapsed_ms = transfer.start_event.elapsed_time(transfer.end_event)
logger.debug(
"EC %s: %d entr%s (%d bytes) took %.3f ms",
direction,
len(transfer.completions),
"y" if len(transfer.completions) == 1 else "ies",
transfer.num_bytes,
elapsed_ms,
)
self._buf_pool.release(transfer.bufs)
self._stream_pool.append(transfer.stream)
self._event_pool.append(transfer.start_event)
self._event_pool.append(transfer.end_event)
return done
def save_caches(
self,
@@ -78,7 +195,14 @@ class ECCPUWorker:
mm_hash: str,
connector_metadata: ECCPUConnectorMetadata,
) -> None:
"""Fill descriptor buffers directly for batched flush."""
"""Fill descriptor buffers directly for batched flush.
Registers `encoder_cache[mm_hash]` against the stream the copy will run
on, so the caching allocator will not hand that memory to another
allocation before the copy has read it. The descriptor buffers hold raw
addresses, and the caller is free to evict the entry as soon as this
step's saves are dispatched.
"""
if not self._is_save_rank:
return
block_ids = connector_metadata.saves.get(mm_hash)
@@ -99,43 +223,95 @@ class ECCPUWorker:
if self._save_bufs is None:
total = sum(len(v) for v in connector_metadata.saves.values())
self._save_bufs = self._buf_pool.acquire(total)
self._save_stream = self._acquire_stream()
self._save_mm_hashes = []
assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel()
bufs = self._save_bufs
src_base = src.view(-1).view(torch.uint8).data_ptr()
assert self._save_stream is not None
src.record_stream(self._save_stream)
# Device addresses are unsigned and can have the high bit set (XPU USM),
# so they are added up in uint64; signed arithmetic rejects them.
src_base = np.uint64(src.view(-1).view(torch.uint8).data_ptr())
dst_base = self._region.blocks.data_ptr()
idx = self._save_count
# Entries are sized by placeholder count, which can exceed the number
# of embeddings, so blocks the output never reaches get no descriptor.
n_used = -(-total_bytes // block_size)
if n_used:
slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used])
run_bytes = num_blocks * block_size
# The output can stop partway into the final run.
run_bytes[-1] = total_bytes - slots[-1] * block_size
bufs.add_copies(
self._save_count,
src_base + slots.astype(np.uint64) * np.uint64(block_size),
dst_base + first_blocks * block_size,
run_bytes,
)
self._save_count += slots.size
for i, block_idx in enumerate(block_ids):
start = i * block_size
bufs.set_ptrs(idx, src_base + start, dst_base + block_idx * block_size)
bufs.sizes[idx] = min(block_size, total_bytes - start)
idx += 1
self._save_count = idx
self._save_bytes += total_bytes
self._save_mm_hashes.append(mm_hash)
def flush_saves(self) -> None:
"""Flush all accumulated saves in a single swap_blocks_batch call."""
"""Flush all accumulated saves in a single swap_blocks_batch call.
Runs the copy on the stream the batch acquired when it opened, gated
behind the compute stream that produced the encoder outputs, brackets it
with start/end events, and enqueues the batch as in-flight. The stream,
events, and descriptor buffers are recycled once the end event fires (see
`_collect_finished`), which is also when the saved mm_hashes become
safe to mark ready.
"""
if self._save_count == 0:
return
bufs = self._save_bufs
assert bufs is not None
stream = self._save_stream
assert bufs is not None and stream is not None
src_ptrs, dst_ptrs, sizes = bufs.src_ptrs, bufs.dst_ptrs, bufs.sizes
n = self._save_count
swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n])
num_bytes = self._save_bytes
# Gate the GPU→CPU copy behind the compute stream: it reads the encoder
# outputs the model just produced there.
stream.wait_stream(current_platform.current_stream())
start_event = self._acquire_event()
end_event = self._acquire_event()
with current_platform.stream(stream):
start_event.record(stream)
swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n])
end_event.record(stream)
self._inflight_saves.append(
Transfer(
start_event=start_event,
end_event=end_event,
completions=self._save_mm_hashes,
bufs=bufs,
stream=stream,
num_bytes=num_bytes,
)
)
self._buf_pool.release(bufs)
self._save_bufs = None
self._save_stream = None
self._save_count = 0
self._save_bytes = 0
self._save_mm_hashes = []
def start_load_caches(
self,
encoder_cache: dict[str, torch.Tensor],
connector_metadata: ECCPUConnectorMetadata,
) -> None:
"""Consumer path: single batched copy of all loads from mmap→GPU."""
"""Consumer path: single batched copy of all loads from mmap→GPU.
The destination is registered against the compute stream that consumes
it, so the caching allocator will not hand that memory to a later load
while reads of it are still queued there.
"""
if not connector_metadata.loads:
return
@@ -145,59 +321,99 @@ class ECCPUWorker:
device_type = current_platform.device_type
src_base = blocks.data_ptr()
# Pre-filter: only hashes not already in encoder_cache.
load_items = {
h: idxs
for h, idxs in connector_metadata.loads.items()
if h not in encoder_cache
}
if not load_items:
return
# Copy every dispatched load. The scheduler routes an input here only
# after missing the GPU encoder cache, so a resident hit is not expected;
# copying unconditionally means every participating rank reports each
# transfer exactly once, which is what the scheduler counts to release
# the pin.
# Ordered by first block so entries that neighbour each other in the
# region land next to each other in the destination and coalesce into
# one descriptor. Every loop below consumes this same order.
load_items = sorted(
connector_metadata.loads.items(), key=lambda kv: kv[1][1][0]
)
total_blocks = sum(len(block_ids) for _, (_, block_ids) in load_items)
total_blocks = sum(len(idxs) for idxs in load_items.values())
with current_platform.stream(self._load_stream):
stream = self._acquire_stream()
compute_stream = current_platform.current_stream()
start_event = self._acquire_event()
end_event = self._acquire_event()
with current_platform.stream(stream):
# Single contiguous destination buffer for all loads.
dst_buf = torch.empty(
total_blocks, block_size, dtype=torch.int8, device=device_type
)
dst_buf_base = dst_buf.data_ptr()
dst_buf.record_stream(compute_stream)
# See save_caches: device addresses are added up in uint64.
dst_buf_base = np.uint64(dst_buf.data_ptr())
bufs = self._buf_pool.acquire(total_blocks)
src_ptrs = bufs.src_ptrs[:total_blocks]
dst_ptrs = bufs.dst_ptrs[:total_blocks]
sizes = bufs.sizes[:total_blocks]
sizes[:] = block_size
op_idx = 0
for block_ids in load_items.values():
for block_idx in block_ids:
bufs.set_ptrs(
op_idx,
src_base + block_idx * block_size,
dst_buf_base + op_idx * block_size,
)
op_idx += 1
slots, first_blocks, num_blocks = _coalesce_runs(
list(chain.from_iterable(ids for _, (_, ids) in load_items))
)
bufs.add_copies(
0,
src_base + first_blocks * block_size,
dst_buf_base + slots.astype(np.uint64) * np.uint64(block_size),
num_blocks * block_size,
)
op_idx = slots.size
src_ptrs = bufs.src_ptrs[:op_idx]
dst_ptrs = bufs.dst_ptrs[:op_idx]
sizes = bufs.sizes[:op_idx]
start_event.record(stream)
swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any=True)
self._buf_pool.release(bufs)
end_event.record(stream)
self._inflight_loads.append(
Transfer(
start_event=start_event,
end_event=end_event,
completions=[tid for _, (tid, _) in load_items],
bufs=bufs,
stream=stream,
num_bytes=total_blocks * block_size,
)
)
# Slice contiguous buffer into per-hash views.
offset = 0
for mm_hash, block_ids in load_items.items():
for mm_hash, (_, block_ids) in load_items:
n = len(block_ids)
encoder_cache[mm_hash] = (
dst_buf[offset : offset + n].view(dtype).reshape(n, -1)
)
offset += n
current_platform.current_stream().wait_stream(self._load_stream)
current_platform.current_stream().wait_stream(stream)
def build_connector_worker_meta(self) -> ECCPUWorkerMetadata | None:
"""Report the GPU copies that completed on this rank this step: saved
mm_hashes and loaded transfer ids.
Returns None when nothing finished, so the scheduler sees no payload.
"""
completed_saves = self._collect_finished(self._inflight_saves, "save")
completed_loads = self._collect_finished(self._inflight_loads, "load")
if not completed_saves and not completed_loads:
return None
return ECCPUWorkerMetadata(
completed_saves=completed_saves,
completed_loads=completed_loads,
)
def shutdown(self) -> None:
self._load_stream.synchronize()
for transfer in (*self._inflight_saves, *self._inflight_loads):
transfer.end_event.synchronize()
self._save_bufs = None
self._save_stream = None
self._save_count = 0
self._save_bytes = 0
self._save_mm_hashes = []
self._inflight_saves.clear()
self._inflight_loads.clear()
self._stream_pool.clear()
self._event_pool.clear()
try:
self._region.cleanup()
except Exception:
vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py
+35 -14
View File
@@ -22,24 +22,43 @@ class DescriptorBuffers(NamedTuple):
src_ptrs: torch.Tensor
dst_ptrs: torch.Tensor
sizes: torch.Tensor
# Numpy aliases of src_ptrs/dst_ptrs, written via set_ptrs().
# Numpy aliases of the three tensors, written via add_copies().
src_np: np.ndarray
dst_np: np.ndarray
sizes_np: np.ndarray
def set_ptrs(self, idx: int, src: int, dst: int) -> None:
"""Record the source and destination address of descriptor *idx*.
def add_copies(
self,
first: int,
sources: np.ndarray,
destinations: np.ndarray,
byte_counts: np.ndarray,
) -> None:
"""Store a batch of copies in these buffers, the first at *first*.
TODO(torch>=2.14): drop this indirection and assign the tensors
directly once the minimum supported torch is 2.14. The numpy detour
exists only because torch's setitem unpacks the value as a signed
long long before 2.14 (pytorch#191458), rejecting XPU USM pointers
>= 2**63, while the two's-complement rewrite those versions accept is
in turn rejected by uint64 from 2.14 on. Numpy casts against the
array dtype and so works on either. `sizes` holds byte counts and
needs no such care.
One copy is one source address, one destination address, and a number
of bytes. `swap_blocks_batch` performs every copy held in the buffers
in a single call, so a caller fills in all of a step's copies before
asking for any of them to run, and passes *first* to say where its own
copies begin.
Whole batches go in at once because storing a single number into a
torch tensor costs microseconds. Adding copies one at a time can take
longer than performing them.
TODO(torch>=2.14): write the tensors directly and drop these numpy
aliases once the minimum supported torch is 2.14. They exist only
because torch's setitem unpacks the value as a signed long long before
2.14 (pytorch#191458), rejecting XPU USM addresses >= 2**63, while the
two's-complement rewrite those versions accept is in turn rejected by
uint64 from 2.14 on. Numpy casts against the array's own type and so
works on either. Byte counts are small enough to need none of this
care, and use their alias only to match the addresses.
"""
self.src_np[idx] = src
self.dst_np[idx] = dst
last = first + len(sources)
self.src_np[first:last] = sources
self.dst_np[first:last] = destinations
self.sizes_np[first:last] = byte_counts
class DescriptorBufferPool:
@@ -62,7 +81,9 @@ class DescriptorBufferPool:
if bufs.src_ptrs.numel() >= n:
return bufs
src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3))
return DescriptorBuffers(src, dst, sizes, src.numpy(), dst.numpy())
return DescriptorBuffers(
src, dst, sizes, src.numpy(), dst.numpy(), sizes.numpy()
)
def release(self, bufs: DescriptorBuffers) -> None:
"""Return a buffer triple to the pool for reuse."""