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.

Functions:

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 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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 
48
49
50
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 
53
54
55
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 
39
40
41
42
43
44
45
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
    )