vLLM/vllm/v1/worker/gpu/eplb_utils.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from collections.abc import Callable, Iterator
from contextlib import contextmanager
from functools import wraps
from typing import TYPE_CHECKING, Any

import torch
import torch.nn as nn

from vllm.config import ModelConfig
from vllm.distributed.eplb.eplb_state import EplbState
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import (
    get_mixture_of_experts_model,
)

if TYPE_CHECKING:
    from vllm.v1.worker.gpu.model_runner import GPUModelRunner

logger = init_logger(__name__)


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


@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])


def _remove_all_requests(model_runner: "GPUModelRunner") -> None:
    for req_id in list(model_runner.req_states.req_id_to_index):
        model_runner._remove_request(req_id)


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

    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 draft_moe_model is None:
            return False

        assert not self.parallel_config.enable_elastic_ep, (
            "Elastic EP is not supported with draft model."
        )
        assert speculative_config is not None
        assert speculative_config.draft_model_config is not None
        assert self.state is not None
        self.state.add_model(draft_moe_model, speculative_config.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)