[Refactor][EC Connector] Add backend extension points to ECCPUWorker (#54033)
Signed-off-by: Akine-Ko <2923601304@qq.com>
Akine-Ko committed
Sep 9, 2026 at 11:13 UTC
41541636fd2075a2a88d9894b5dc3b05c396ed28
2 files changed
+117
-9
tests/v1/ec_connector/unit/cpu/worker/test_worker.py
+71
-1
@@ -38,7 +38,10 @@ from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
38
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
39
ECSharedRegion,
40
)
41
-from vllm.distributed.ec_transfer.ec_connector.cpu.worker import ECCPUWorker
41
+from vllm.distributed.ec_transfer.ec_connector.cpu.worker import (
42
+ ECCPUTransferDirection,
43
+ ECCPUWorker,
44
+)
45
from vllm.platforms import current_platform
46
47
# ── shape constants ──────────────────────────────────────────────────────────
@@ -199,6 +202,71 @@ def _warm_up_stream_pools(worker: ECCPUWorker) -> None:
202
assert worker._stream_pool, "warm-up did not recycle a stream"
203
204
205
+# ── backend extension points ────────────────────────────────────────────────
206
+
207
+
208
+def test_save_path_delegates_transfer_submission(make_worker):
209
+ worker = make_worker()
210
+ stream = Mock()
211
+ worker._acquire_stream = Mock(return_value=stream)
212
+ worker._acquire_event = Mock(side_effect=[Mock(), Mock()])
213
+ worker._submit_transfer = Mock()
214
+
215
+ src = MagicMock()
216
+ src.numel.return_value = _HIDDEN_DIM
217
+ src.element_size.return_value = _DTYPE.itemsize
218
+ src.view.return_value = src
219
+ src.data_ptr.return_value = 1000
220
+
221
+ platform = Mock()
222
+ platform.current_stream.return_value = Mock()
223
+ platform.stream.return_value = contextlib.nullcontext()
224
+ with patch(
225
+ "vllm.distributed.ec_transfer.ec_connector.cpu.worker.current_platform",
226
+ platform,
227
+ ):
228
+ worker.save_caches({"h": src}, "h", _meta(saves={"h": [3]}))
229
+ worker.flush_saves()
230
+
231
+ worker._submit_transfer.assert_called_once()
232
+ _, count, direction = worker._submit_transfer.call_args.args
233
+ assert count == 1
234
+ assert direction == ECCPUTransferDirection.DEVICE_TO_HOST
235
+
236
+
237
+def test_load_path_delegates_transfer_submission(make_worker):
238
+ worker = make_worker()
239
+ stream = Mock()
240
+ compute_stream = Mock()
241
+ worker._acquire_stream = Mock(return_value=stream)
242
+ worker._acquire_event = Mock(side_effect=[Mock(), Mock()])
243
+ worker._submit_transfer = Mock()
244
+
245
+ dst_buf = MagicMock()
246
+ dst_buf.data_ptr.return_value = 1000
247
+
248
+ platform = Mock()
249
+ platform.device_type = "cpu"
250
+ platform.current_stream.return_value = compute_stream
251
+ platform.stream.return_value = contextlib.nullcontext()
252
+ with (
253
+ patch(
254
+ "vllm.distributed.ec_transfer.ec_connector.cpu.worker.current_platform",
255
+ platform,
256
+ ),
257
+ patch(
258
+ "vllm.distributed.ec_transfer.ec_connector.cpu.worker.torch.empty",
259
+ return_value=dst_buf,
260
+ ),
261
+ ):
262
+ worker.start_load_caches({}, _meta(loads={"h": [3]}))
263
+
264
+ worker._submit_transfer.assert_called_once()
265
+ _, count, direction = worker._submit_transfer.call_args.args
266
+ assert count == 1
267
+ assert direction == ECCPUTransferDirection.HOST_TO_DEVICE
268
+
269
+
270
# ── save_caches ──────────────────────────────────────────────────────────────
271
272
@@ -792,8 +860,10 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):
860
worker._inflight_loads = deque()
861
worker._stream_pool = []
862
worker._event_pool = []
863
+ worker._shutdown_transfer_backend = MagicMock()
864
865
worker.shutdown()
866
+ worker._shutdown_transfer_backend.assert_called_once()
867
mock_region.cleanup.assert_called_once()
868
869
mock_region.cleanup.side_effect = RuntimeError("boom")
vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py
+46
-8
@@ -10,6 +10,7 @@ blocks to copy in each direction.
10
from collections import deque
11
from collections.abc import Sequence
12
from dataclasses import dataclass
13
+from enum import Enum, auto
14
from itertools import chain
15
from typing import TYPE_CHECKING
16
@@ -40,6 +41,11 @@ if TYPE_CHECKING:
41
logger = init_logger(__name__)
42
43
44
+class ECCPUTransferDirection(Enum):
45
+ HOST_TO_DEVICE = auto()
46
+ DEVICE_TO_HOST = auto()
47
+
48
+
49
def _coalesce_runs(
50
block_ids: Sequence[int],
51
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -111,8 +117,7 @@ class ECCPUWorker:
117
# Model dtype; used to reinterpret raw int8 blocks on load.
118
self._dtype = vllm_config.model_config.dtype
119
114
- if is_pin_memory_available():
115
- self._region.pin_memory()
120
+ self._pin_shared_region()
121
122
# All TP/PCP ranks hold identical encoder output. Only one rank
123
# per mmap needs to write — saves host memory bandwidth.
@@ -188,6 +193,34 @@ class ECCPUWorker:
193
self._event_pool.append(transfer.end_event)
194
return done
195
196
+ def _pin_shared_region(self) -> None:
197
+ """Prepare the shared region for device transfers."""
198
+ if is_pin_memory_available():
199
+ self._region.pin_memory()
200
+
201
+ def _submit_transfer(
202
+ self,
203
+ bufs: DescriptorBuffers,
204
+ count: int,
205
+ direction: ECCPUTransferDirection,
206
+ ) -> None:
207
+ """Submit a transfer on the current backend stream."""
208
+ src_ptrs = bufs.src_ptrs[:count]
209
+ dst_ptrs = bufs.dst_ptrs[:count]
210
+ sizes = bufs.sizes[:count]
211
+ if direction == ECCPUTransferDirection.HOST_TO_DEVICE:
212
+ swap_blocks_batch(
213
+ src_ptrs,
214
+ dst_ptrs,
215
+ sizes,
216
+ is_src_access_order_any=True,
217
+ )
218
+ else:
219
+ swap_blocks_batch(src_ptrs, dst_ptrs, sizes)
220
+
221
+ def _shutdown_transfer_backend(self) -> None:
222
+ """Release backend-specific transfer resources."""
223
+
224
def save_caches(
225
self,
226
encoder_cache: dict[str, torch.Tensor],
@@ -269,7 +302,6 @@ class ECCPUWorker:
302
bufs = self._save_bufs
303
stream = self._save_stream
304
assert bufs is not None and stream is not None
272
- src_ptrs, dst_ptrs, sizes = bufs.src_ptrs, bufs.dst_ptrs, bufs.sizes
305
n = self._save_count
306
num_bytes = self._save_bytes
307
@@ -280,7 +312,11 @@ class ECCPUWorker:
312
end_event = self._acquire_event()
313
with current_platform.stream(stream):
314
start_event.record(stream)
283
- swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n])
315
+ self._submit_transfer(
316
+ bufs,
317
+ n,
318
+ ECCPUTransferDirection.DEVICE_TO_HOST,
319
+ )
320
end_event.record(stream)
321
322
self._inflight_saves.append(
@@ -357,12 +393,13 @@ class ECCPUWorker:
393
num_blocks * block_size,
394
)
395
op_idx = slots.size
360
- src_ptrs = bufs.src_ptrs[:op_idx]
361
- dst_ptrs = bufs.dst_ptrs[:op_idx]
362
- sizes = bufs.sizes[:op_idx]
396
397
start_event.record(stream)
365
- swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any=True)
398
+ self._submit_transfer(
399
+ bufs,
400
+ op_idx,
401
+ ECCPUTransferDirection.HOST_TO_DEVICE,
402
+ )
403
end_event.record(stream)
404
self._inflight_loads.append(
405
Transfer(
@@ -404,6 +441,7 @@ class ECCPUWorker:
441
def shutdown(self) -> None:
442
for transfer in (*self._inflight_saves, *self._inflight_loads):
443
transfer.end_event.synchronize()
444
+ self._shutdown_transfer_backend()
445
self._save_bufs = None
446
self._save_stream = None
447
self._save_count = 0