| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared types for the ECCPUConnector scheduler and worker delegates.""" |
| 4 | |
| 5 | from dataclasses import dataclass, field |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.distributed.ec_transfer.ec_connector.base import ( |
| 11 | ECConnectorMetadata, |
| 12 | ECConnectorWorkerMetadata, |
| 13 | ) |
| 14 | from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( |
| 15 | ECSharedRegion, |
| 16 | ) |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from vllm.config import VllmConfig |
| 20 | |
| 21 | |
| 22 | @dataclass |
| 23 | class ECCPUConnectorMetadata(ECConnectorMetadata): |
| 24 | """Per-step scheduler → worker payload for the ECCPUConnector. |
| 25 | |
| 26 | Populated by `ECCPUScheduler.build_connector_meta`; consumed by |
| 27 | `ECCPUWorker` via the mixin's `bind_connector_metadata`. |
| 28 | """ |
| 29 | |
| 30 | # Producer role: mm_hashes the scheduler has just allocated CPU |
| 31 | # blocks for this step; the worker's save_caches copies |
| 32 | # encoder_cache[mm_hash] → mmap at these block IDs. |
| 33 | saves: dict[str, list[int]] = field(default_factory=dict) |
| 34 | |
| 35 | # Consumer role: mm_hashes whose bytes are available in the local mmap, |
| 36 | # mapped to (transfer_id, block_ids); the worker's start_load_caches |
| 37 | # copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id |
| 38 | # once its copy lands. The id is minted per dispatch, so two loads of the |
| 39 | # same mm_hash in different steps are distinct and a late report from the |
| 40 | # earlier one cannot release the later one's pin. |
| 41 | loads: dict[str, tuple[int, list[int]]] = field(default_factory=dict) |
| 42 | |
| 43 | |
| 44 | @dataclass |
| 45 | class ECCPUWorkerMetadata(ECConnectorWorkerMetadata): |
| 46 | """Per-step worker → scheduler payload for the ECCPUConnector. |
| 47 | |
| 48 | Reports the GPU copies that completed this step: saved mm_hashes become |
| 49 | safe to mark ready, and loaded transfers become safe to unpin once every |
| 50 | participating rank has reported them. Built by |
| 51 | `ECCPUWorker.build_connector_worker_meta`; consumed by |
| 52 | `ECCPUScheduler.update_connector_output`. |
| 53 | |
| 54 | Loads are reported by transfer id rather than mm_hash because every rank |
| 55 | copies the same blocks: `aggregate` concatenates, so an id appears once |
| 56 | per reporting rank and the scheduler can count participants off the list. |
| 57 | """ |
| 58 | |
| 59 | completed_saves: list[str] = field(default_factory=list) |
| 60 | completed_loads: list[int] = field(default_factory=list) |
| 61 | |
| 62 | def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetadata: |
| 63 | assert isinstance(other, ECCPUWorkerMetadata) |
| 64 | self.completed_saves.extend(other.completed_saves) |
| 65 | self.completed_loads.extend(other.completed_loads) |
| 66 | return self |
| 67 | |
| 68 | |
| 69 | def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int: |
| 70 | """Return the per-token hidden dimension for encoder cache entries. |
| 71 | |
| 72 | For most models this equals the LLM's hidden size. Qwen3-VL (and any |
| 73 | future model with deepstack visual encoding) is an exception: the ViT |
| 74 | concatenates its own output with features from N decoder layers before |
| 75 | storing in encoder_cache, producing a tensor of width |
| 76 | ``out_hidden_size * (1 + N)`` per visual token. Using the plain LLM |
| 77 | hidden size would under-allocate EC blocks and silently truncate the |
| 78 | transferred data, leading to a shape mismatch on the consumer. |
| 79 | """ |
| 80 | model_config = vllm_config.model_config |
| 81 | hf_config = getattr(model_config, "hf_config", None) |
| 82 | vision_config = ( |
| 83 | getattr(hf_config, "vision_config", None) if hf_config is not None else None |
| 84 | ) |
| 85 | if vision_config is not None: |
| 86 | out_hidden_size = getattr(vision_config, "out_hidden_size", None) |
| 87 | deepstack_indexes = getattr(vision_config, "deepstack_visual_indexes", None) |
| 88 | if out_hidden_size is not None and deepstack_indexes: |
| 89 | return out_hidden_size * (1 + len(deepstack_indexes)) |
| 90 | return model_config.get_inputs_embeds_size() |
| 91 | |
| 92 | |
| 93 | def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: |
| 94 | """Build the EC mmap region from `vllm_config`. |
| 95 | |
| 96 | Both `ECCPUScheduler` and `ECCPUWorker` call this to get the same |
| 97 | shared region (same engine_id, same block_size_bytes). |
| 98 | """ |
| 99 | ec_config = vllm_config.ec_transfer_config |
| 100 | assert ec_config is not None, "ec_transfer_config required to build region" |
| 101 | |
| 102 | dp_index = vllm_config.parallel_config.data_parallel_index |
| 103 | engine_id = f"{vllm_config.instance_id}_dp{dp_index}" |
| 104 | |
| 105 | dtype = vllm_config.model_config.dtype |
| 106 | hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) |
| 107 | element_size = torch.empty(0, dtype=dtype).element_size() |
| 108 | block_size_bytes = hidden_dim * element_size |
| 109 | |
| 110 | cpu_bytes = ec_config.ec_connector_extra_config.get("ec_cpu_bytes") |
| 111 | if not cpu_bytes: |
| 112 | raise ValueError("ec_cpu_bytes must be specified in ec_connector_extra_config") |
| 113 | cpu_bytes = int(cpu_bytes) |
| 114 | num_blocks = cpu_bytes // block_size_bytes |
| 115 | |
| 116 | return ECSharedRegion( |
| 117 | engine_id=engine_id, |
| 118 | num_blocks=num_blocks, |
| 119 | block_size_bytes=block_size_bytes, |
| 120 | ) |