| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable, Iterator |
| 5 | from contextlib import contextmanager |
| 6 | from functools import wraps |
| 7 | from typing import TYPE_CHECKING, Any |
| 8 | |
| 9 | import torch |
| 10 | import torch.nn as nn |
| 11 | |
| 12 | from vllm.config import ModelConfig |
| 13 | from vllm.distributed.eplb.eplb_state import EplbState |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.model_executor.models.interfaces import ( |
| 16 | get_mixture_of_experts_model, |
| 17 | ) |
| 18 | |
| 19 | if TYPE_CHECKING: |
| 20 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | |
| 25 | def step_eplb_after(*, is_dummy: bool = False) -> Callable: |
| 26 | """Step EPLB after a model runner method completes successfully.""" |
| 27 | |
| 28 | def decorator(fn: Callable) -> Callable: |
| 29 | @wraps(fn) |
| 30 | def wrapper(self: Any, *args, **kwargs) -> Any: |
| 31 | result = fn(self, *args, **kwargs) |
| 32 | if kwargs.get("skip_eplb", False): |
| 33 | return result |
| 34 | |
| 35 | is_profile = kwargs.get("is_profile", False) if is_dummy else False |
| 36 | self.eplb.step(is_dummy=is_dummy, is_profile=is_profile) |
| 37 | return result |
| 38 | |
| 39 | return wrapper |
| 40 | |
| 41 | return decorator |
| 42 | |
| 43 | |
| 44 | @contextmanager |
| 45 | def preserve_serving_state(model_runner: "GPUModelRunner") -> Iterator[None]: |
| 46 | """Keep the elastic EP warmup out of the request pool and the KV cache.""" |
| 47 | # After the drain only parked streaming sessions still hold a slot, and |
| 48 | # those are re-added from NewRequestData on their next chunk. |
| 49 | _remove_all_requests(model_runner) |
| 50 | model_runner.block_tables.redirect_writes_to_null_block = True |
| 51 | try: |
| 52 | yield |
| 53 | finally: |
| 54 | _remove_all_requests(model_runner) |
| 55 | model_runner.block_tables.redirect_writes_to_null_block = False |
| 56 | if model_runner.kv_block_zeroer is not None: |
| 57 | model_runner.kv_block_zeroer.zero_block_ids([0]) |
| 58 | |
| 59 | |
| 60 | def _remove_all_requests(model_runner: "GPUModelRunner") -> None: |
| 61 | for req_id in list(model_runner.req_states.req_id_to_index): |
| 62 | model_runner._remove_request(req_id) |
| 63 | |
| 64 | |
| 65 | class EPLBController: |
| 66 | def __init__(self, parallel_config: Any, device: torch.device): |
| 67 | self.parallel_config = parallel_config |
| 68 | self.device = device |
| 69 | self.state: EplbState | None = None |
| 70 | self.suppressed = False |
| 71 | self._has_registered_models = False |
| 72 | |
| 73 | def prepare_load(self) -> None: |
| 74 | self.state = None |
| 75 | self._has_registered_models = False |
| 76 | if self.parallel_config.enable_eplb: |
| 77 | self.state = EplbState(self.parallel_config, self.device) |
| 78 | |
| 79 | def maybe_register_speculator( |
| 80 | self, |
| 81 | speculator: Any | None, |
| 82 | speculative_config: Any | None, |
| 83 | load_dummy_weights: bool, |
| 84 | ) -> bool: |
| 85 | # if speculator is a moe model, add it to eplb |
| 86 | if ( |
| 87 | speculator is None |
| 88 | or not hasattr(speculator, "model") |
| 89 | or not self.parallel_config.enable_eplb |
| 90 | or load_dummy_weights |
| 91 | ): |
| 92 | return False |
| 93 | |
| 94 | draft_model = speculator.model |
| 95 | draft_moe_model = get_mixture_of_experts_model(draft_model) |
| 96 | if draft_moe_model is None: |
| 97 | return False |
| 98 | |
| 99 | assert not self.parallel_config.enable_elastic_ep, ( |
| 100 | "Elastic EP is not supported with draft model." |
| 101 | ) |
| 102 | assert speculative_config is not None |
| 103 | assert speculative_config.draft_model_config is not None |
| 104 | assert self.state is not None |
| 105 | self.state.add_model(draft_moe_model, speculative_config.draft_model_config) |
| 106 | speculator.set_eplb_state(self.state) |
| 107 | self._has_registered_models = True |
| 108 | return True |
| 109 | |
| 110 | def maybe_register_model(self, model: nn.Module, model_config: ModelConfig) -> bool: |
| 111 | if not self.parallel_config.enable_eplb: |
| 112 | return False |
| 113 | |
| 114 | moe_model = get_mixture_of_experts_model(model) |
| 115 | if moe_model is None: |
| 116 | return False |
| 117 | |
| 118 | logger.info_once( |
| 119 | "EPLB is enabled for MoE part of model %s.", model_config.model |
| 120 | ) |
| 121 | assert self.state is not None |
| 122 | self.state.add_model(moe_model, model_config) |
| 123 | self._has_registered_models = True |
| 124 | return True |
| 125 | |
| 126 | def maybe_start_async_loop( |
| 127 | self, eplb_models_added: bool, load_dummy_weights: bool = False |
| 128 | ) -> None: |
| 129 | # A scaling-up worker registers to join the communicator handshake, but |
| 130 | # its groups are not live yet, so _perform_eplb_reshuffle() starts this. |
| 131 | if load_dummy_weights: |
| 132 | return |
| 133 | if eplb_models_added and self.state is not None and self.state.is_async: |
| 134 | self.state.start_async_loop() |
| 135 | |
| 136 | def step(self, is_dummy: bool = False, is_profile: bool = False) -> None: |
| 137 | if ( |
| 138 | not self.parallel_config.enable_eplb |
| 139 | or self.suppressed |
| 140 | or self.state is None |
| 141 | or not self._has_registered_models |
| 142 | ): |
| 143 | return |
| 144 | |
| 145 | log_stats = self.parallel_config.eplb_config.log_balancedness |
| 146 | self.state.step(is_dummy, is_profile, log_stats=log_stats) |
| 147 | |
| 148 | def prepare_forward( |
| 149 | self, |
| 150 | model_config: ModelConfig, |
| 151 | num_unpadded_tokens: int, |
| 152 | ubatch_slices: list | None = None, |
| 153 | ) -> None: |
| 154 | if self.state is None or not self.parallel_config.enable_eplb: |
| 155 | return |
| 156 | self.state.prepare_forward(model_config, num_unpadded_tokens, ubatch_slices) |
| 157 | |
| 158 | def setup_from_mapping( |
| 159 | self, |
| 160 | model_config: ModelConfig, |
| 161 | expanded_physical_to_logical: torch.Tensor, |
| 162 | ) -> None: |
| 163 | # Update in place: a fresh EplbState would build a second communicator |
| 164 | # over a group the other ranks have already left. |
| 165 | assert self.state is not None |
| 166 | self.state.update_mapping(model_config, expanded_physical_to_logical) |