diff --git a/docs/features/ec_cpu_connector.md b/docs/features/ec_cpu_connector.md index 2a6df25f6d..18274b6dd9 100644 --- a/docs/features/ec_cpu_connector.md +++ b/docs/features/ec_cpu_connector.md @@ -96,7 +96,7 @@ EC transfer is configured via `--ec-transfer-config` (CLI) or the `ec_transfer_c | --- | --- | --- | --- | | `ec_enable_nixl` | `bool` | No (default `false`) | Enables NIXL P2P transfer in addition to local CPU offload. Omitted or `false` imports no NIXL/ZMQ. Extra config is not type coerced, so a string value is parsed: `"true"`, `"1"`, `"yes"` enable it, anything else does not. | | `consumer_ack_timeout_s` | `float` | No (default `2.0`) | How long a consumer waits for an `XferAck` before giving up on a read. The producer answers `XferReq`s from its own scheduler step, so its reply latency scales with the encoder's `--max-num-batched-tokens`: a loaded encoder whose steps run longer than this makes consumers abandon reads the producer is about to grant. Raise it for large encoder batches. | -| `ec_cpu_bytes` | `int` | Yes | Total size, in bytes, of the shared CPU mmap region. `ECCPUConnector` raises `ValueError` if unset. Block count = `ec_cpu_bytes // block_size_bytes`, where `block_size_bytes = hidden_dim * dtype.element_size()` (`hidden_dim` accounts for Qwen3-VL deepstack: `out_hidden_size * (1 + num_deepstack_layers)`). | +| `ec_cpu_bytes` | `int` | Yes | Total size, in bytes, of the shared CPU mmap region. `ECCPUConnector` raises `ValueError` if unset. Block count = `ec_cpu_bytes // block_size_bytes`, where each block holds one LLM input-embedding vector. Entries allocate enough blocks for their modality-specific output shape, including visual DeepStack features. | ### Environment Variables diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py index 24ea717cab..ef86505057 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py @@ -22,12 +22,16 @@ class _Pos: self.offset = offset self.length = length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1, identifier=None): self.mm_hash = mm_hash self.identifier = identifier if identifier is not None else mm_hash self.mm_position = _Pos(0, length) + self.modality = "image" class _Request: @@ -80,7 +84,7 @@ def _make_scheduler( 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()] + return [transfer_id for transfer_id, _, _ in meta.loads.values()] def _load_blocks(meta, mm_hash: str) -> list[int]: @@ -90,7 +94,7 @@ def _load_blocks(meta, mm_hash: str) -> list[int]: def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int): """Pre-populate a ready cache entry backed by real blocks.""" - s._cache.alloc(mm_hash, n_blocks) + s._cache.alloc(mm_hash, n_blocks, (n_blocks, _BS // s._dtype.itemsize)) s._cache.mark_ready(mm_hash) diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index c168970147..0b64d68ca6 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -97,10 +97,9 @@ def test_second_instance_opens_existing_file_and_shares_memory(): r1.cleanup() -def test_dense_dp_engines_get_distinct_regions(monkeypatch): +def test_dense_dp_engines_get_distinct_regions(): """Dense DP engines reset data_parallel_rank to 0 but share instance_id, so the region must be keyed by data_parallel_index.""" - monkeypatch.setattr(common, "_get_encoder_cache_hidden_dim", lambda cfg: 8) instance_id = str(uuid.uuid4()) regions = [] try: @@ -113,6 +112,7 @@ def test_dense_dp_engines_get_distinct_regions(monkeypatch): cfg.instance_id = instance_id cfg.parallel_config = parallel_config cfg.model_config.dtype = torch.float16 + cfg.model_config.get_inputs_embeds_size.return_value = 8 cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64} regions.append(common.create_ec_shared_region(cfg)) diff --git a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py index f4f5f03f69..4ff0979e30 100644 --- a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py +++ b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py @@ -91,7 +91,9 @@ def _make_region() -> ECSharedRegion: def _vllm_config(rank: int = 0) -> Mock: - return create_ec_vllm_config(rank=rank, dtype=_DTYPE) + config = create_ec_vllm_config(rank=rank, dtype=_DTYPE) + config.model_config.get_inputs_embeds_size.return_value = _HIDDEN_DIM + return config def _meta( @@ -105,7 +107,7 @@ def _meta( return ECCPUConnectorMetadata( saves=saves or {}, loads={ - mm_hash: (idx, block_ids) + mm_hash: (idx, block_ids, (len(block_ids), _HIDDEN_DIM)) for idx, (mm_hash, block_ids) in enumerate((loads or {}).items()) }, ) @@ -904,21 +906,29 @@ 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) - scheduler = ECCPUScheduler(_vllm_config()) + config = _vllm_config() + # DeepStack-wide output: each embedding spans four blocks. + width = 4 * _HIDDEN_DIM + config.ec_transfer_config.mm_encoder_output_widths = {"image": width} + scheduler = ECCPUScheduler(config) # -- Step 1: scheduler allocates, worker saves -- - n_blocks = 3 - src = torch.arange( - n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE - ).reshape(n_blocks, _HIDDEN_DIM) + n_tokens = 2 + src = torch.arange(n_tokens * width, dtype=_DTYPE, device=DEVICE_TYPE).reshape( + n_tokens, width + ) class _Pos: offset = 0 - length = n_blocks + length = n_tokens + + def get_num_embeds(self): + return self.length class _Feature: identifier = "img_001" mm_position = _Pos() + modality = "image" class _Request: request_id = "req_e2e" @@ -948,8 +958,9 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch): scheduler.update_state_after_alloc(_Request(), 0) meta_load = scheduler.build_connector_meta(scheduler_output=None) assert "img_001" in meta_load.loads - _, load_blocks = meta_load.loads["img_001"] + _, load_blocks, shape = meta_load.loads["img_001"] assert load_blocks == meta_save.saves["img_001"] + assert shape == tuple(src.shape) load_cache: dict[str, torch.Tensor] = {} worker.start_load_caches(load_cache, meta_load) diff --git a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py index 26155d8b6e..20dd32416c 100644 --- a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py +++ b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py @@ -45,6 +45,34 @@ KEY = ( SPEC = TensorSpec((2, 2), "torch.float32", 16) +def test_push_shapes_follow_each_modality(monkeypatch): + """Each push is sized at its modality's measured encoder output width.""" + from vllm.distributed.ec_transfer.ec_connector.mooncake import scheduler + + monkeypatch.setattr(scheduler, "ensure_mooncake_available", lambda: None) + widths = {"image": 8192, "audio": 2048} + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths=widths) + request = SimpleNamespace( + request_id="mixed", + mm_features=[SimpleNamespace(identifier=m, modality=m) for m in widths], + get_num_encoder_embeds=lambda index: 2, + ec_transfer_params={ + "consumer_zmq": "tcp://consumer:1234", + "ec_items": [{"mm_hash": m, "transfer_id": m} for m in widths], + }, + ) + instance = scheduler.ECMooncakeScheduler(config) + try: + for index in range(len(widths)): + instance.update_state_after_alloc(request, index) + metadata = instance.build_connector_meta( + SimpleNamespace(free_encoder_mm_hashes=[], preempted_req_ids=set()) + ) + assert [spec.shape for spec in metadata.pushes] == [(2, 8192), (2, 2048)] + finally: + instance.close() + + def test_embedding_key_combines_namespace_and_identifier(): config = create_ec_vllm_config(ec_role="ec_producer") config.ec_transfer_config.ec_connector_extra_config.update( diff --git a/tests/v1/ec_connector/unit/test_encoder_output_width.py b/tests/v1/ec_connector/unit/test_encoder_output_width.py new file mode 100644 index 0000000000..b7d7c1a822 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_encoder_output_width.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup measurement of the encoder output width EC producers size by.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config +from vllm.distributed.ec_transfer.ec_connector import encoder_output_width as eow + +pytestmark = pytest.mark.cpu_test + + +class _DeepStackModel: + """Emits each modality at its own width, as Qwen3-Omni does.""" + + widths = {"image": 4 * 16, "audio": 16} + + def embed_multimodal(self, *, modality): + return [torch.zeros(3, self.widths[modality])] + + +class _Budget: + def __init__(self, *args, **kwargs): + self.mm_max_toks_per_item = dict.fromkeys(_DeepStackModel.widths, 8) + + def get_dummy_encoder_profile_inputs(self, modality, max_items_per_batch): + assert max_items_per_batch == 1 + return [(modality, None)] + + +def _worker(first_pp_rank=True, supports_mm=True): + runner = SimpleNamespace( + supports_mm_inputs=supports_mm, + get_model=_DeepStackModel, + device=torch.device("cpu"), + ) + pp_group = SimpleNamespace(is_first_rank=first_pp_rank) + return SimpleNamespace(model_runner=runner, vllm_config=None), pp_group + + +@pytest.fixture +def patched(monkeypatch): + import vllm.distributed.parallel_state as ps + import vllm.multimodal.encoder_budget as budget_mod + import vllm.multimodal.utils as mm_utils + + monkeypatch.setattr(budget_mod, "MultiModalBudget", _Budget) + monkeypatch.setattr( + mm_utils, + "group_and_batch_mm_kwargs", + lambda items, **_: iter([(items[0][0], 1, {"modality": items[0][0]})]), + ) + return lambda group: monkeypatch.setattr(ps, "get_pp_group", lambda: group) + + +def test_worker_reports_the_encoder_output_width_per_modality(patched): + """The width comes from the tensor the encoder emits, not from config.""" + worker, group = _worker() + patched(group) + assert eow._measure_on_worker(worker) == {"image": 64, "audio": 16} + + +@pytest.mark.parametrize("first_pp_rank, supports_mm", [(False, True), (True, False)]) +def test_ranks_without_an_encoder_report_nothing(patched, first_pp_rank, supports_mm): + worker, group = _worker(first_pp_rank, supports_mm) + patched(group) + assert eow._measure_on_worker(worker) == {} + + +def test_producer_records_the_widths_its_encoder_ranks_agree_on(): + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths={}) + reports = [{"image": 64}, {}, {"image": 64}] + eow.measure_encoder_output_widths(config, lambda fn: reports) + assert config.ec_transfer_config.mm_encoder_output_widths == {"image": 64} + assert eow.get_encoder_output_width(config, "image") == 64 + + +@pytest.mark.parametrize( + "reports, match", + [ + ([{}, {}], "could not measure"), + ([{"image": 64}, {"image": 16}], "disagree"), + ], + ids=["no-encoder", "disagree"], +) +def test_producer_refuses_to_start_without_a_trustworthy_width(reports, match): + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths={}) + with pytest.raises(ValueError, match=match): + eow.measure_encoder_output_widths(config, lambda fn: reports) diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py index a3e87f9b78..8c9e14563f 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -25,6 +25,9 @@ class _Pos: def __init__(self, offset, length): self.offset, self.length = offset, length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1, data=None): @@ -85,6 +88,9 @@ class _FakeTransport: pass +_MISSING = object() + + def _consumer_sched(monkeypatch): def _region(cfg): return ECSharedRegion( @@ -95,16 +101,15 @@ def _consumer_sched(monkeypatch): # ec_enable_nixl defaults to False, so this builds gate-off; the fields # below flip on the NIXL consumer state directly. - s = ECCPUScheduler(create_ec_vllm_config(ec_role="ec_consumer")) + # No measured widths: a consumer takes each shape from the announcement. + s = ECCPUScheduler( + create_ec_vllm_config(ec_role="ec_consumer", encoder_output_widths={}) + ) # Turn on NIXL consumer state without constructing real transports. s._nixl_enabled = True s._transport = _FakeTransport() s._data = None s._compat_hash = "c" - # _setup_nixl normally computes these from model_config; set them - # directly since this helper builds gate-off then flips fields on. - s._hidden_dim = _HID - s._element_size = _ES s._metadata_resolver._cache["image"] = {"image_grid_thw"} return s @@ -115,24 +120,32 @@ def _params(mm_hash, length): "peer_host": "h", "peer_port": 1, "size_bytes": length * _HID * _ES, + "shape": [length, _HID], } } def test_new_remote_read_defers_then_completes(monkeypatch): + """The read is sized from the announced shape, not a local width.""" + width = 4 * _HID s = _consumer_sched(monkeypatch) fake = _FakeSession() # Route _start_xfer to our fake session instead of real ZMQ/NIXL, but # still reserve a real not-ready cache entry so mark_ready works. - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + assert size == width * _ES + assert shape == (1, width) + entry = s._cache.alloc(mm_hash, size // _BS, shape) assert entry is not None fake.started.append(mm_hash) return True monkeypatch.setattr(s, "_start_xfer", _fake_start) - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) + params = _params("h1", 1) + params["h1"]["size_bytes"] = width * _ES + params["h1"]["shape"] = [1, width] + req = _Request([_Feature("h1", 1)], params=params) # Step 1: unseen remote item -> read started, request deferred. assert s.ensure_cache_available(req, 0) is False @@ -155,6 +168,7 @@ def test_new_remote_read_defers_then_completes(monkeypatch): s.update_state_after_alloc(req, 0) meta = s.build_connector_meta(scheduler_output=None) assert "h1" in meta.loads + assert meta.loads["h1"][2] == (1, width) assert "h1" not in s._step_completed # cleared by promote # Step 3: still cached -> admitted directly, no new transfer. @@ -176,8 +190,8 @@ def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload): s = _consumer_sched(monkeypatch) fake = _FakeSession() - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + entry = s._cache.alloc(mm_hash, 1, shape) assert entry is not None fake.started.append(mm_hash) return True @@ -230,8 +244,8 @@ def test_remote_wait_budget_survives_retries_and_long_steps(monkeypatch, retry): fake = _FakeSession() s._sessions[("h", 1)] = fake - def start(mm_hash, info, size): - assert s._cache.alloc(mm_hash, 1) is not None + def start(mm_hash, info, size, shape): + assert s._cache.alloc(mm_hash, 1, shape) is not None fake.started.append(mm_hash) return True @@ -261,8 +275,8 @@ def test_retryable_read_re_requests_without_admitting(monkeypatch): s = _consumer_sched(monkeypatch) fake = _FakeSession() - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + entry = s._cache.alloc(mm_hash, 1, shape) assert entry is not None fake.started.append(mm_hash) return True @@ -287,14 +301,29 @@ def test_retryable_read_re_requests_without_admitting(monkeypatch): @pytest.mark.parametrize( - "size_fields", - [{"size_bytes": 999}, {"size_bytes": None}, {"size_bytes": "big"}, {}], - ids=["mismatch", "null-size", "text-size", "no-size"], + "fields", + [ + {"size_bytes": 999}, + {"size_bytes": None}, + {"size_bytes": "big"}, + {"size_bytes": _MISSING}, + {"shape": _MISSING}, + {"shape": [2, _HID]}, + ], + ids=[ + "size-mismatch", + "null-size", + "text-size", + "no-size", + "no-shape", + "wrong-rows", + ], ) -def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): - """Invalid remote sizes fail the request without raising in the scheduler.""" +def test_invalid_announcement_fails_the_request(monkeypatch, fields): + """An unusable announced size or shape fails the request, not the engine.""" s = _consumer_sched(monkeypatch) - announced = {"peer_host": "h", "peer_port": 1, **size_fields} + announced = {**_params("h1", 1)["h1"], **fields} + announced = {k: v for k, v in announced.items() if v is not _MISSING} req = _Request([_Feature("h1", 1)], params={"h1": announced}) assert s.ensure_cache_available(req, 0) is False assert s.take_unavailable_requests() == {"r1"} @@ -306,7 +335,7 @@ def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): def test_deferral_budget_is_per_request(monkeypatch): """Shared hashes have independent wait budgets and cancellation cleanup.""" s = _consumer_sched(monkeypatch) - monkeypatch.setattr(s._cache, "alloc", lambda key, n: None) + monkeypatch.setattr(s._cache, "alloc", lambda key, n, shape=None: None) params = _params("h1", 1) old = _Request([_Feature("h1", 1)], params=params, req_id="old") new = _Request([_Feature("h1", 1)], params=params, req_id="new") @@ -419,7 +448,7 @@ def test_in_flight_hash_defers_without_second_transfer(monkeypatch): s = _consumer_sched(monkeypatch) started: list[str] = [] - def _spy_start_xfer(mm_hash, info, size): + def _spy_start_xfer(mm_hash, info, size, shape): started.append(mm_hash) return True diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py index d7ef19b5f4..167f333f73 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -16,6 +16,9 @@ class _Pos: def __init__(self, offset, length): self.offset, self.length = offset, length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1): @@ -53,11 +56,8 @@ def test_request_finished_producer_emits_params(monkeypatch): # Simulate NIXL-enabled producer bookkeeping without building real NIXL. s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - # _setup_nixl normally computes these from model_config; set them - # directly since this test builds gate-off then flips fields on. - s._hidden_dim, s._element_size = 32, 2 # feature length=2, hidden_dim=32, element_size=2 -> 128 bytes -> 2 blocks. - entry = s._cache.alloc("h1", 2) + entry = s._cache.alloc("h1", 2, (2, 32)) assert entry is not None s._cache.mark_ready("h1") @@ -71,6 +71,7 @@ def test_request_finished_producer_emits_params(monkeypatch): "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, + "shape": [2, 32], } } s.shutdown() @@ -84,8 +85,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._hidden_dim, s._element_size = 32, 2 - s._cache.alloc("h1", 2) # allocated but not marked ready + s._cache.alloc("h1", 2, (2, 32)) # allocated but not marked ready delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) assert delay is False @@ -96,6 +96,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, + "shape": [2, 32], } } s.shutdown() @@ -105,7 +106,6 @@ def test_request_finished_skips_unallocated_entry(monkeypatch): s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._hidden_dim, s._element_size = 32, 2 # No alloc() for "h1" at all — e.g. the cache was full at save time. delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) diff --git a/tests/v1/ec_connector/unit/utils.py b/tests/v1/ec_connector/unit/utils.py index a718bc63b6..705a8ff7ad 100644 --- a/tests/v1/ec_connector/unit/utils.py +++ b/tests/v1/ec_connector/unit/utils.py @@ -18,12 +18,16 @@ _EXECUTOR_BACKEND = "mp" _PCP_SIZE = 1 +_EMBED_WIDTH = 32 + + def create_ec_vllm_config( *, ec_role: ECRole = "ec_both", tensor_parallel_size: int = 1, rank: int = 0, dtype: torch.dtype = torch.float16, + encoder_output_widths: dict[str, int] | None = None, ) -> Mock: """Build a `VllmConfig` stand-in for EC connector unit tests. @@ -31,10 +35,8 @@ def create_ec_vllm_config( 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. + `model_config` supplies a dtype and embedding width without resolving an + HF model or inspecting its architecture. Args: ec_role: EC role this instance plays. `is_ec_producer` and @@ -42,6 +44,8 @@ def create_ec_vllm_config( tensor_parallel_size: TP degree. rank: Global rank of this worker. dtype: Encoder cache dtype. + encoder_output_widths: Per-modality encoder output width, as a + producer measures it at startup. Defaults to the embedding width. Returns: A `VllmConfig`-specced mock carrying the configs above. @@ -59,9 +63,15 @@ def create_ec_vllm_config( ec_role=ec_role, engine_id=str(uuid.uuid4()), ) + ec_transfer_config.mm_encoder_output_widths = ( + dict.fromkeys(("image", "video", "audio"), _EMBED_WIDTH) + if encoder_output_widths is None + else encoder_output_widths + ) model_config = Mock(spec=ModelConfig) model_config.dtype = dtype + model_config.get_inputs_embeds_size.return_value = _EMBED_WIDTH vllm_config = Mock(spec=VllmConfig) vllm_config.ec_transfer_config = ec_transfer_config diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index 5c2b8e589c..b3ca10a5a6 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -55,6 +55,11 @@ class ECTransferConfig: """The Python module path to dynamically load the EC connector from. Only supported in V1.""" + mm_encoder_output_widths: dict[str, int] = field(default_factory=dict, init=False) + """Per-token width of the encoder output for each modality, measured by EC + producers at startup. It can exceed the LM embedding width (e.g. DeepStack). + """ + def compute_hash(self) -> str: """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py index eff16e08b0..461cd6a55a 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py @@ -33,12 +33,14 @@ class ECCPUConnectorMetadata(ECConnectorMetadata): saves: 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 + # mapped to (transfer_id, block_ids, shape); 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) + loads: dict[str, tuple[int, list[int], tuple[int, int]]] = field( + default_factory=dict + ) @dataclass @@ -66,30 +68,6 @@ class ECCPUWorkerMetadata(ECConnectorWorkerMetadata): return self -def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int: - """Return the per-token hidden dimension for encoder cache entries. - - For most models this equals the LLM's hidden size. Qwen3-VL (and any - future model with deepstack visual encoding) is an exception: the ViT - concatenates its own output with features from N decoder layers before - storing in encoder_cache, producing a tensor of width - ``out_hidden_size * (1 + N)`` per visual token. Using the plain LLM - hidden size would under-allocate EC blocks and silently truncate the - transferred data, leading to a shape mismatch on the consumer. - """ - model_config = vllm_config.model_config - hf_config = getattr(model_config, "hf_config", None) - vision_config = ( - getattr(hf_config, "vision_config", None) if hf_config is not None else None - ) - if vision_config is not None: - out_hidden_size = getattr(vision_config, "out_hidden_size", None) - deepstack_indexes = getattr(vision_config, "deepstack_visual_indexes", None) - if out_hidden_size is not None and deepstack_indexes: - return out_hidden_size * (1 + len(deepstack_indexes)) - return model_config.get_inputs_embeds_size() - - def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: """Build the EC mmap region from `vllm_config`. @@ -103,7 +81,7 @@ def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: engine_id = f"{vllm_config.instance_id}_dp{dp_index}" dtype = vllm_config.model_config.dtype - hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) + hidden_dim = vllm_config.model_config.get_inputs_embeds_size() element_size = torch.empty(0, dtype=dtype).element_size() block_size_bytes = hidden_dim * element_size diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py b/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py index 85bd3cae32..e20bf48f0b 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py @@ -14,9 +14,9 @@ import json import msgspec -# Bump when the on-wire shape (XferReq/XferAck) changes in a +# Bump when the wire format or tensor/block layout changes in a # backward-incompatible way. Peers reject mismatched versions. -EC_CONNECTOR_VERSION: int = 1 +EC_CONNECTOR_VERSION: int = 2 class XferStatus(enum.IntEnum): diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py index 65403980b9..567f2d59c9 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -16,12 +16,14 @@ import torch from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUConnectorMetadata, ECCPUWorkerMetadata, - _get_encoder_cache_hidden_dim, create_ec_shared_region, ) from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( EmbeddingCache, ) +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + get_encoder_output_width, +) from vllm.distributed.ec_transfer.ec_connector.utils import ( PlaceholderMetadataResolver, collect_ec_item_metadata, @@ -49,6 +51,7 @@ class ECCPUScheduler: """Scheduler delegate for the ECCPUConnector.""" def __init__(self, vllm_config: "VllmConfig") -> None: + self._vllm_config = vllm_config ec_config = vllm_config.ec_transfer_config assert ec_config is not None self._is_producer: bool = ec_config.is_ec_producer @@ -61,8 +64,8 @@ class ECCPUScheduler: # mm_hash → block IDs allocated this step for GPU→mmap saves. self._pending_saves: 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]]] = {} + # mm_hash → (transfer_id, block IDs, shape) to load from mmap→GPU. + self._pending_loads: dict[str, tuple[int, list[int], tuple[int, int]]] = {} # Dispatched loads awaiting completion reports, keyed by transfer id: # transfer_id → (mm_hash, reports still outstanding). The pin taken at @@ -108,11 +111,9 @@ class ECCPUScheduler: self._unrecoverable: set[str] = set() self._peer_host: str | None = None self._peer_port: int | None = None - # Model shape for size checks + compat hash; only set by - # _setup_nixl (the gate-off path never touches model_config). - self._dtype: torch.dtype | None = None - self._hidden_dim: int = 0 - self._element_size: int = 0 + dtype = vllm_config.model_config.dtype + assert isinstance(dtype, torch.dtype) + self._dtype = dtype self._ack_timeout_s: float = 0.0 if self._nixl_enabled: try: @@ -148,9 +149,6 @@ class ECCPUScheduler: ) engine_id = self._ec_config.engine_id assert engine_id is not None - self._dtype = vllm_config.model_config.dtype - self._hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) - self._element_size = torch.empty(0, dtype=self._dtype).element_size() # How long a consumer waits for an XferAck. The producer answers from # its scheduler step, so its reply latency scales with the encoder's # batch size: a deployment whose steps run longer must raise this. @@ -296,22 +294,30 @@ class ECCPUScheduler: # held by a quarantined/settling DMA and cannot be reused. return False - expected = pos.length * self._hidden_dim * self._element_size try: size = int(remote["size_bytes"]) + shape = tuple(int(dim) for dim in remote["shape"]) except (KeyError, TypeError, ValueError) as error: raise _RemoteUnavailable("the announced size was unusable") from error - if size != expected: + # The width is the producer's to report; the row count is not. + num_embeds = pos.get_num_embeds() + if ( + len(shape) != 2 + or shape[0] != num_embeds + or size != self._num_blocks(shape) * self._region.block_size_bytes + ): logger.warning( - "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", + "EC consumer: size mismatch mm_hash=%s announced=%d shape=%s " + "expected %d rows", mm_hash, size, - expected, + shape, + num_embeds, ) raise _RemoteUnavailable("the announced size was wrong") try: - started = self._start_xfer(mm_hash, remote, expected) + started = self._start_xfer(mm_hash, remote, size, (shape[0], shape[1])) except Exception as error: logger.exception( "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash @@ -377,7 +383,11 @@ class ECCPUScheduler: return failed def _start_xfer( - self, mm_hash: str, info: "dict[str, Any]", size_bytes: int + self, + mm_hash: str, + info: "dict[str, Any]", + size_bytes: int, + shape: tuple[int, int], ) -> bool: """Allocate a not-ready cache entry and start a NIXL READ into it. @@ -391,7 +401,7 @@ class ECCPUScheduler: ) n_blocks = max(1, ceil(size_bytes / self._region.block_size_bytes)) - entry = self._cache.alloc(mm_hash, n_blocks) + entry = self._cache.alloc(mm_hash, n_blocks, shape) if entry is None: logger.debug( "EC consumer: cache full for mm_hash=%s (%d blocks); deferring", @@ -491,22 +501,36 @@ class ECCPUScheduler: """ self._step_completed.clear() + def _num_blocks(self, shape: tuple[int, ...]) -> int: + size_bytes = shape[0] * shape[1] * self._dtype.itemsize + block_size = self._region.block_size_bytes + return (size_bytes + block_size - 1) // block_size + def update_state_after_alloc(self, request: "Request", index: int) -> None: feature = request.mm_features[index] mm_hash = feature.identifier if self._is_producer and self._cache.get(mm_hash) is None: - entry = self._cache.alloc(mm_hash, feature.mm_position.length) + shape = ( + feature.mm_position.get_num_embeds(), + get_encoder_output_width(self._vllm_config, feature.modality), + ) + entry = self._cache.alloc(mm_hash, self._num_blocks(shape), shape) if entry is not None: self._pending_saves[mm_hash] = list(entry.block_ids) 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: + assert entry.shape is not None self._cache.pin(mm_hash) transfer_id = self._next_transfer_id self._next_transfer_id += 1 - self._pending_loads[mm_hash] = (transfer_id, list(entry.block_ids)) + self._pending_loads[mm_hash] = ( + transfer_id, + list(entry.block_ids), + entry.shape, + ) self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports) def build_connector_meta( @@ -608,13 +632,13 @@ class ECCPUScheduler: # A read arriving before the save lands is NACKed NACK_NOT_READY, # which the consumer retries on a later step rather than treating # as a miss. - size_bytes = ( - feature.mm_position.length * self._hidden_dim * self._element_size - ) + assert entry.shape is not None + size_bytes = len(entry.block_ids) * self._region.block_size_bytes items[mm_hash].update( peer_host=self._peer_host, peer_port=self._peer_port, size_bytes=size_bytes, + shape=list(entry.shape), ) logger.debug( "EC producer: announcing NIXL-readable encodings req_id=%s items=%s", diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py index 2fd4b01059..046fba0574 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py @@ -14,13 +14,20 @@ from collections import OrderedDict class CacheEntry: - """A single cache entry. Read `.block_ids` and `.ready` freely; - mutations only through EmbeddingCache methods.""" + """A single cache entry. Read `.block_ids`, `.shape` and `.ready` freely; + mutations only through EmbeddingCache methods. - __slots__ = ("block_ids", "_pin_count") + `shape` is the encoder output the blocks hold, so a load can restore it + without knowing how wide this model's encoder output is. + """ - def __init__(self, block_ids: tuple[int, ...]) -> None: + __slots__ = ("block_ids", "shape", "_pin_count") + + def __init__( + self, block_ids: tuple[int, ...], shape: tuple[int, int] | None = None + ) -> None: self.block_ids = block_ids + self.shape = shape self._pin_count = -1 # not ready @property @@ -76,7 +83,9 @@ class EmbeddingCache: with self._lock: return len(self._entries) != len(self._entries_free_list) - def alloc(self, key: str, n_blocks: int) -> CacheEntry | None: + def alloc( + self, key: str, n_blocks: int, shape: tuple[int, int] | None = None + ) -> CacheEntry | None: """Allocate *n_blocks* for *key*, evicting as needed. The entry starts not-ready. Returns None if there is not enough @@ -97,7 +106,7 @@ class EmbeddingCache: # 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) + entry = CacheEntry(block_ids, shape) self._entries[key] = entry return entry diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py index ef7ff0f751..95c6c96128 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py @@ -267,8 +267,7 @@ class ECCPUWorker: # 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() - # Entries are sized by placeholder count, which can exceed the number - # of embeddings, so blocks the output never reaches get no descriptor. + # Copy only the tensor bytes, excluding any block padding. n_used = -(-total_bytes // block_size) if n_used: slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used]) @@ -367,7 +366,7 @@ class ECCPUWorker: 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(block_ids) for _, (_, block_ids, _) in load_items) stream = self._acquire_stream() compute_stream = current_platform.current_stream() @@ -384,7 +383,7 @@ class ECCPUWorker: bufs = self._buf_pool.acquire(total_blocks) slots, first_blocks, num_blocks = _coalesce_runs( - list(chain.from_iterable(ids for _, (_, ids) in load_items)) + list(chain.from_iterable(ids for _, (_, ids, _) in load_items)) ) bufs.add_copies( 0, @@ -405,7 +404,7 @@ class ECCPUWorker: Transfer( start_event=start_event, end_event=end_event, - completions=[tid for _, (tid, _) in load_items], + completions=[tid for _, (tid, _, _) in load_items], bufs=bufs, stream=stream, num_bytes=total_blocks * block_size, @@ -414,11 +413,10 @@ class ECCPUWorker: # Slice contiguous buffer into per-hash views. offset = 0 - for mm_hash, (_, block_ids) in load_items: + for mm_hash, (_, block_ids, shape) in load_items: n = len(block_ids) - encoder_cache[mm_hash] = ( - dst_buf[offset : offset + n].view(dtype).reshape(n, -1) - ) + flat = dst_buf[offset : offset + n].view(dtype).reshape(-1) + encoder_cache[mm_hash] = flat[: shape[0] * shape[1]].view(shape) offset += n current_platform.current_stream().wait_stream(stream) diff --git a/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py b/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py new file mode 100644 index 0000000000..c1a270bb2b --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder output widths for sizing EC transfers. + +An encoder can emit features wider than the LM embedding (DeepStack packs +several layers on the last dim) by an amount only the model knows. Producers +size transfers before the first real item is encoded, so the engine measures +the width once at startup by encoding a dummy item per modality. +""" + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +@torch.inference_mode() +def _measure_on_worker(worker: Any) -> dict[str, int]: + """Encode one dummy item per modality; runs on every worker. + + Returns {} on ranks that hold no encoder. + """ + from vllm.distributed.parallel_state import get_pp_group + from vllm.multimodal import MULTIMODAL_REGISTRY + from vllm.multimodal.encoder_budget import MultiModalBudget + from vllm.multimodal.utils import group_and_batch_mm_kwargs + from vllm.utils.torch_utils import PIN_MEMORY + + runner = worker.model_runner + if not runner.supports_mm_inputs or not get_pp_group().is_first_rank: + return {} + budget = MultiModalBudget( + worker.vllm_config, MULTIMODAL_REGISTRY, enable_cache=False + ) + model = runner.get_model() + widths: dict[str, int] = {} + for modality in budget.mm_max_toks_per_item: + items = budget.get_dummy_encoder_profile_inputs(modality, 1) + if not items: + continue + _, _, mm_kwargs = next( + group_and_batch_mm_kwargs( + items, device=runner.device, pin_memory=PIN_MEMORY + ) + ) + outputs = model.embed_multimodal(**mm_kwargs) + widths[modality] = int(outputs[0].shape[-1]) + del outputs + return widths + + +def measure_encoder_output_widths( + vllm_config: "VllmConfig", collective_rpc: Callable[..., list[Any]] +) -> None: + """Record the encoder output widths on an EC producer's config. + + Call after the model is loaded and before the KV cache is allocated: the + Scheduler's connector reads the widths, and once the KV cache holds the + memory budget there is no room left for a dummy encode. + + Raises: + ValueError: No worker could encode a dummy item, or workers disagree. + + """ + ec_config = vllm_config.ec_transfer_config + if ec_config is None or not ec_config.is_ec_producer: + return + reports = [widths for widths in collective_rpc(_measure_on_worker) if widths] + if not reports: + raise ValueError( + "EC producer could not measure the encoder output width: the model " + "encoded no dummy multimodal input." + ) + if any(widths != reports[0] for widths in reports): + raise ValueError(f"Workers disagree on the encoder output widths: {reports}") + logger.info("EC producer encoder output widths: %s", reports[0]) + ec_config.mm_encoder_output_widths = reports[0] + + +def get_encoder_output_width(vllm_config: "VllmConfig", modality: str) -> int: + """Encoder output width for `modality`, as measured at producer startup. + + Raises: + ValueError: No width was measured for `modality`, so a transfer + cannot be sized. + + """ + assert vllm_config.ec_transfer_config is not None + widths = vllm_config.ec_transfer_config.mm_encoder_output_widths + if modality not in widths: + raise ValueError( + f"No encoder output width was measured for modality {modality!r}; " + f"measured: {sorted(widths)}" + ) + return widths[modality] diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py index 2a607118db..597ba767df 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py @@ -20,8 +20,8 @@ import torch from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorMetadata, ) -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - _get_encoder_cache_hidden_dim, +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + get_encoder_output_width, ) from vllm.distributed.ec_transfer.ec_connector.mooncake.config import ( _RESERVATION_TTL_SECONDS, @@ -85,9 +85,7 @@ class ECMooncakeScheduler: self._is_consumer = config.is_consumer self._control_addr = config.control_addr self._push_wait_timeout = config.push_wait_timeout_s - self._encoder_cache_hidden_dim = ( - _get_encoder_cache_hidden_dim(vllm_config) if config.is_producer else None - ) + self._vllm_config = vllm_config self._model_config = vllm_config.model_config self._control_client = ControlClient(config.control_timeout_ms) self._control_executor = ThreadPoolExecutor( @@ -447,8 +445,10 @@ class ECMooncakeScheduler: def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec: dtype = self._model_config.dtype assert isinstance(dtype, torch.dtype) - assert self._encoder_cache_hidden_dim is not None - shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim) + hidden_dim = get_encoder_output_width( + self._vllm_config, request.mm_features[index].modality + ) + shape = (request.get_num_encoder_embeds(index), hidden_dim) return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize) def update_state_after_alloc(self, request: Any, index: int) -> None: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index bbca64a58a..35cc8cc186 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -27,6 +27,9 @@ from vllm.distributed import ( cleanup_dist_env_and_memory, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + measure_encoder_output_widths, +) from vllm.envs import enable_envs_cache from vllm.logger import configure_logging, init_logger from vllm.logging_utils.dump_input import dump_engine_exception @@ -148,6 +151,8 @@ class EngineCore: if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: self._eep_scale_up_before_kv_init() + measure_encoder_output_widths(vllm_config, self.collective_rpc) + # Setup KV Caches and update CacheConfig after profiling. kv_cache_config = self._initialize_kv_caches(vllm_config) self.structured_output_manager = StructuredOutputManager(vllm_config)