Skip to content

vllm.model_executor.model_loader.weight_cache.utils

Helpers shared by the weight cache daemon, the IPC loader and the engine.

These live outside protocol so callers that only need to know whether a draft is cached, or how a daemon group is named, do not have to import the wire format.

Classes:

Functions:

ExternalWeightCharges dataclass

Track a worker's external weight charges and those freed by disk fallback.

Attributes:

  • current_bytes (int) –

    External bytes held right now: charges minus fallback removals.

Source code in vllm/model_executor/model_loader/weight_cache/utils.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
@dataclass
class ExternalWeightCharges:
    """Track a worker's external weight charges and those freed by disk fallback."""

    fallback_bytes: int = field(default=0, init=False)
    _charges: dict[tuple[str, bool], int] = field(default_factory=dict, init=False)

    @property
    def current_bytes(self) -> int:
        """External bytes held right now: charges minus fallback removals."""
        return sum(self._charges.values())

    @contextmanager
    def track(self) -> Generator[None, None, None]:
        token = _active_external_weight_charges.set(self)
        try:
            yield
        finally:
            _active_external_weight_charges.reset(token)

current_bytes property

External bytes held right now: charges minus fallback removals.

build_warmup_runner(vllm_config, local_rank, *, is_draft, model_config)

Create a model runner for the daemon's warmup dummy runs.

Source code in vllm/model_executor/model_loader/weight_cache/utils.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def build_warmup_runner(
    vllm_config: VllmConfig,
    local_rank: int,
    *,
    is_draft: bool,
    model_config: "ModelConfig",
) -> "GPUModelRunner":
    """Create a model runner for the daemon's warmup dummy runs."""
    if is_draft:
        vllm_config = replace(
            vllm_config, model_config=model_config, speculative_config=None
        )
    device = torch.device(current_platform.device_type, local_rank)
    with set_current_vllm_config(vllm_config):
        if vllm_config.use_v2_model_runner:
            from vllm.v1.worker.gpu.model_runner import GPUModelRunner as Runner
        else:
            from vllm.v1.worker.gpu_model_runner import GPUModelRunner as Runner
        runner = Runner(vllm_config, device)
    # The two runner classes share the warmup surface; the nominal type
    # stays V1's GPUModelRunner.
    return cast("GPUModelRunner", runner)

format_daemon_role(is_draft)

Name of a daemon group: the target model or the draft.

Source code in vllm/model_executor/model_loader/weight_cache/utils.py 
91
92
93
def format_daemon_role(is_draft: bool) -> str:
    """Name of a daemon group: the target model or the draft."""
    return "draft" if is_draft else "target"

format_socket_role_suffix(is_draft)

Socket-name suffix keeping the draft group distinct from the target.

Source code in vllm/model_executor/model_loader/weight_cache/utils.py 
96
97
98
def format_socket_role_suffix(is_draft: bool) -> str:
    """Socket-name suffix keeping the draft group distinct from the target."""
    return "_draft" if is_draft else ""

is_draft_model_cacheable(speculative_config)

Whether the daemon serves the speculative draft as a separate role.

Source code in vllm/model_executor/model_loader/weight_cache/utils.py 
82
83
84
85
86
87
88
def is_draft_model_cacheable(speculative_config: SpeculativeConfig | None) -> bool:
    """Whether the daemon serves the speculative draft as a separate role."""
    return (
        speculative_config is not None
        and speculative_config.method in WEIGHT_CACHE_DRAFT_METHODS
        and speculative_config.draft_model_config is not None
    )