Skip to content

vllm.v1.worker.gpu.eplb_utils

Functions:

draft_model_supports_eplb(speculative_config, draft_moe_model)

Return whether a draft MoE model should register with EPLB.

Source code in vllm/v1/worker/gpu/eplb_utils.py 
32
33
34
35
36
37
38
39
40
41
42
43
def draft_model_supports_eplb(
    speculative_config: Any | None,
    draft_moe_model: MixtureOfExperts | None,
) -> bool:
    """Return whether a draft MoE model should register with EPLB."""
    if draft_moe_model is None:
        return False
    if speculative_config is None or speculative_config.draft_model_config is None:
        return False
    if getattr(speculative_config, "method", None) == "dspark":
        return dspark_draft_supports_eplb(speculative_config.draft_model_config)
    return True

eplb_draft_model_name(model_config)

Return the EPLB step-log display name for a draft model.

Source code in vllm/v1/worker/gpu/eplb_utils.py 
27
28
29
def eplb_draft_model_name(model_config: ModelConfig) -> str:
    """Return the EPLB step-log display name for a draft model."""
    return f"{model_config.model} (draft)"

preserve_serving_state(model_runner)

Keep the elastic EP warmup out of the request pool and the KV cache.

Source code in vllm/v1/worker/gpu/eplb_utils.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
@contextmanager
def preserve_serving_state(model_runner: "GPUModelRunner") -> Iterator[None]:
    """Keep the elastic EP warmup out of the request pool and the KV cache."""
    # After the drain only parked streaming sessions still hold a slot, and
    # those are re-added from NewRequestData on their next chunk.
    _remove_all_requests(model_runner)
    model_runner.block_tables.redirect_writes_to_null_block = True
    try:
        yield
    finally:
        _remove_all_requests(model_runner)
        model_runner.block_tables.redirect_writes_to_null_block = False
        if model_runner.kv_block_zeroer is not None:
            model_runner.kv_block_zeroer.zero_block_ids([0])

step_eplb_after(*, is_dummy=False)

Step EPLB after a model runner method completes successfully.

Source code in vllm/v1/worker/gpu/eplb_utils.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def step_eplb_after(*, is_dummy: bool = False) -> Callable:
    """Step EPLB after a model runner method completes successfully."""

    def decorator(fn: Callable) -> Callable:
        @wraps(fn)
        def wrapper(self: Any, *args, **kwargs) -> Any:
            result = fn(self, *args, **kwargs)
            if kwargs.get("skip_eplb", False):
                return result

            is_profile = kwargs.get("is_profile", False) if is_dummy else False
            self.eplb.step(is_dummy=is_dummy, is_profile=is_profile)
            return result

        return wrapper

    return decorator