| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Helpers shared by the weight cache daemon, the IPC loader and the engine. |
| 4 | |
| 5 | These live outside ``protocol`` so callers that only need to know whether a |
| 6 | draft is cached, or how a daemon group is named, do not have to import the |
| 7 | wire format. |
| 8 | """ |
| 9 | |
| 10 | from typing import Any |
| 11 | |
| 12 | from vllm.config import SpeculativeConfig |
| 13 | from vllm.model_executor.models.interfaces import SupportsEagleBase |
| 14 | |
| 15 | # Speculative methods whose draft model the daemon caches in its own group. |
| 16 | # Other drafts keep loading from disk in the engine. |
| 17 | WEIGHT_CACHE_DRAFT_METHODS = frozenset({"mtp", "eagle", "eagle3"}) |
| 18 | |
| 19 | # Python-side flags that weight loading sets on EAGLE-style drafts; the engine |
| 20 | # never runs load_weights for cached models, so the daemon ships them. |
| 21 | EXPORTED_MODEL_ATTRS = tuple(SupportsEagleBase.__annotations__) |
| 22 | |
| 23 | |
| 24 | def export_model_attrs(model: Any) -> dict[str, bool]: |
| 25 | return { |
| 26 | name: bool(getattr(model, name)) |
| 27 | for name in EXPORTED_MODEL_ATTRS |
| 28 | if hasattr(model, name) |
| 29 | } |
| 30 | |
| 31 | |
| 32 | def is_draft_model_cacheable(speculative_config: SpeculativeConfig | None) -> bool: |
| 33 | """Whether the daemon serves the speculative draft as a separate role.""" |
| 34 | return ( |
| 35 | speculative_config is not None |
| 36 | and speculative_config.method in WEIGHT_CACHE_DRAFT_METHODS |
| 37 | and speculative_config.draft_model_config is not None |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | def format_daemon_role(is_draft: bool) -> str: |
| 42 | """Name of a daemon group: the target model or the draft.""" |
| 43 | return "draft" if is_draft else "target" |
| 44 | |
| 45 | |
| 46 | def format_socket_role_suffix(is_draft: bool) -> str: |
| 47 | """Socket-name suffix keeping the draft group distinct from the target.""" |
| 48 | return "_draft" if is_draft else "" |