[Bugfix][EC Connector] Size transfers by measured encoder output widths (#57904)

Signed-off-by: Tianyu Guo <guoty@inferact.ai>

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
 

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)
 
 

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

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)

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(

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)

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
 

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

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

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

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
 

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

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

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
 

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)

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]

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:

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)