vllm.v1.worker.gpu.eplb_utils

Classes:

• EPLBController –

Functions:

• draft_model_supports_eplb –

  Return whether a draft MoE model should register with EPLB.

• eplb_draft_model_name –

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

• preserve_serving_state –

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

• step_eplb_after –

  Step EPLB after a model runner method completes successfully.

EPLBController

Methods:

• suppress –

  Keep warmup forwards from stepping EPLB or recording expert load.

Source code in vllm/v1/worker/gpu/eplb_utils.py

class EPLBController:
    def __init__(self, parallel_config: Any, device: torch.device):
        self.parallel_config = parallel_config
        self.device = device
        self.state: EplbState | None = None
        self.suppressed = False
        self._has_registered_models = False

    @contextmanager
    def suppress(self) -> Iterator[None]:
        """Keep warmup forwards from stepping EPLB or recording expert load."""
        should_record = (
            self.state.should_record_tensor if self.state is not None else None
        )
        was_suppressed = self.suppressed
        was_recording = should_record.clone() if should_record is not None else None
        self.suppressed = True
        if should_record is not None:
            should_record.fill_(False)
        try:
            yield
        finally:
            if should_record is not None and was_recording is not None:
                should_record.copy_(was_recording)
            self.suppressed = was_suppressed

    def prepare_load(self) -> None:
        self.state = None
        self._has_registered_models = False
        if self.parallel_config.enable_eplb:
            self.state = EplbState(self.parallel_config, self.device)

    def maybe_register_speculator(
        self,
        speculator: Any | None,
        speculative_config: Any | None,
        load_dummy_weights: bool,
    ) -> bool:
        # if speculator is a moe model, add it to eplb
        if (
            speculator is None
            or not hasattr(speculator, "model")
            or not self.parallel_config.enable_eplb
            or load_dummy_weights
        ):
            return False

        draft_model = speculator.model
        draft_moe_model = get_mixture_of_experts_model(draft_model)
        if not draft_model_supports_eplb(speculative_config, draft_moe_model):
            return False
        assert draft_moe_model is not None

        assert not self.parallel_config.enable_elastic_ep, (
            "Elastic EP is not supported with draft model."
        )
        assert speculative_config is not None
        draft_model_config = speculative_config.draft_model_config
        assert draft_model_config is not None
        logger.info_once(
            "EPLB is enabled for MoE part of drafter model %s.",
            draft_model_config.model,
        )
        assert self.state is not None
        self.state.add_model(
            draft_moe_model,
            draft_model_config,
            model_name=eplb_draft_model_name(draft_model_config),
        )
        speculator.set_eplb_state(self.state)
        self._has_registered_models = True
        return True

    def maybe_register_model(self, model: nn.Module, model_config: ModelConfig) -> bool:
        if not self.parallel_config.enable_eplb:
            return False

        moe_model = get_mixture_of_experts_model(model)
        if moe_model is None:
            return False

        logger.info_once(
            "EPLB is enabled for MoE part of model %s.", model_config.model
        )
        assert self.state is not None
        self.state.add_model(moe_model, model_config)
        self._has_registered_models = True
        return True

    def maybe_start_async_loop(
        self, eplb_models_added: bool, load_dummy_weights: bool = False
    ) -> None:
        # A scaling-up worker registers to join the communicator handshake, but
        # its groups are not live yet, so _perform_eplb_reshuffle() starts this.
        if load_dummy_weights:
            return
        if eplb_models_added and self.state is not None and self.state.is_async:
            self.state.start_async_loop()

    def step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
        if (
            not self.parallel_config.enable_eplb
            or self.suppressed
            or self.state is None
            or not self._has_registered_models
        ):
            return

        log_stats = self.parallel_config.eplb_config.log_balancedness
        self.state.step(is_dummy, is_profile, log_stats=log_stats)

    def prepare_forward(
        self,
        model_config: ModelConfig,
        num_unpadded_tokens: int,
        ubatch_slices: list | None = None,
    ) -> None:
        if self.state is None or not self.parallel_config.enable_eplb:
            return
        self.state.prepare_forward(model_config, num_unpadded_tokens, ubatch_slices)

    def setup_from_mapping(
        self,
        model_config: ModelConfig,
        expanded_physical_to_logical: torch.Tensor,
    ) -> None:
        # Update in place: a fresh EplbState would build a second communicator
        # over a group the other ranks have already left.
        assert self.state is not None
        self.state.update_mapping(model_config, expanded_physical_to_logical)

suppress()

Keep warmup forwards from stepping EPLB or recording expert load.

Source code in vllm/v1/worker/gpu/eplb_utils.py

@contextmanager
def suppress(self) -> Iterator[None]:
    """Keep warmup forwards from stepping EPLB or recording expert load."""
    should_record = (
        self.state.should_record_tensor if self.state is not None else None
    )
    was_suppressed = self.suppressed
    was_recording = should_record.clone() if should_record is not None else None
    self.suppressed = True
    if should_record is not None:
        should_record.fill_(False)
    try:
        yield
    finally:
        if should_record is not None and was_recording is not None:
            should_record.copy_(was_recording)
        self.suppressed = was_suppressed

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

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

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

@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

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:
            if kwargs.get("skip_eplb", False):
                # The router kernel records load before step() is reached.
                with self.eplb.suppress():
                    return fn(self, *args, **kwargs)

            result = fn(self, *args, **kwargs)
            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