| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | from typing import Any |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import CompilationConfig, CompilationMode |
| 11 | from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry |
| 12 | from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT |
| 13 | from vllm.v1.worker.gpu import eplb_utils as eplb |
| 14 | from vllm.v1.worker.gpu import model_runner as mrv2 |
| 15 | |
| 16 | |
| 17 | class FakeMemoryProfiler: |
| 18 | def __enter__(self): |
| 19 | self.consumed_memory = 0 |
| 20 | return self |
| 21 | |
| 22 | def __exit__(self, exc_type, exc, tb): |
| 23 | return False |
| 24 | |
| 25 | |
| 26 | class FakeEplbState: |
| 27 | instances: list["FakeEplbState"] = [] |
| 28 | |
| 29 | def __init__(self, parallel_config: Any, device: torch.device): |
| 30 | self.parallel_config = parallel_config |
| 31 | self.device = device |
| 32 | self.add_model_calls: list[tuple[Any, Any]] = [] |
| 33 | self.step_calls: list[tuple[bool, bool, bool]] = [] |
| 34 | self.async_started = False |
| 35 | self.is_async = True |
| 36 | self.update_mapping_calls: list[tuple[Any, torch.Tensor]] = [] |
| 37 | FakeEplbState.instances.append(self) |
| 38 | |
| 39 | def add_model(self, model: Any, model_config: Any) -> None: |
| 40 | self.add_model_calls.append((model, model_config)) |
| 41 | |
| 42 | def step(self, is_dummy: bool, is_profile: bool, *, log_stats: bool) -> None: |
| 43 | self.step_calls.append((is_dummy, is_profile, log_stats)) |
| 44 | |
| 45 | def start_async_loop(self) -> None: |
| 46 | self.async_started = True |
| 47 | |
| 48 | def update_mapping(self, model_config: Any, mapping: torch.Tensor) -> None: |
| 49 | self.update_mapping_calls.append((model_config, mapping)) |
| 50 | |
| 51 | |
| 52 | def _make_runner(**overrides: Any) -> Any: |
| 53 | runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) |
| 54 | runner.device = torch.device("cpu") |
| 55 | runner.model_config = SimpleNamespace(model="test-model", logits_processors=None) |
| 56 | runner.req_states = SimpleNamespace( |
| 57 | device=runner.device, |
| 58 | max_num_reqs=8, |
| 59 | vocab_size=0, |
| 60 | all_token_ids=None, |
| 61 | prompt_len=None, |
| 62 | prefill_len=None, |
| 63 | total_len=None, |
| 64 | ) |
| 65 | runner.load_config = SimpleNamespace(load_format="hf") |
| 66 | runner.parallel_config = SimpleNamespace( |
| 67 | enable_eplb=True, |
| 68 | enable_elastic_ep=False, |
| 69 | enable_batch_sharded_sampling=False, |
| 70 | eplb_config=SimpleNamespace(log_balancedness=True), |
| 71 | ) |
| 72 | runner.compilation_config = CompilationConfig(mode=CompilationMode.NONE) |
| 73 | runner.vllm_config = SimpleNamespace( |
| 74 | compilation_config=runner.compilation_config, |
| 75 | load_config=runner.load_config, |
| 76 | model_config=runner.model_config, |
| 77 | ) |
| 78 | runner.lora_config = None |
| 79 | runner.use_aux_hidden_state_outputs = False |
| 80 | runner.speculative_config = None |
| 81 | runner.speculator = None |
| 82 | runner.pcp_manager = None |
| 83 | runner.num_speculative_steps = 0 |
| 84 | runner.encoder_cache = None |
| 85 | runner.is_pooling_model = False |
| 86 | runner.is_last_pp_rank = True |
| 87 | runner.is_first_pp_rank = True |
| 88 | runner.max_num_reqs = 8 |
| 89 | runner.max_num_tokens = 16 |
| 90 | runner.decode_query_len = 1 |
| 91 | runner.kv_connector = SimpleNamespace( |
| 92 | set_disabled=lambda *_: None, |
| 93 | post_forward=lambda *_, **__: None, |
| 94 | ) |
| 95 | runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device) |
| 96 | runner.jit_warmup_registry = JitWarmupRegistry(runner.vllm_config) |
| 97 | runner.pooling_runner = None |
| 98 | runner.execute_model_state = None |
| 99 | for key, value in overrides.items(): |
| 100 | setattr(runner, key, value) |
| 101 | return runner |
| 102 | |
| 103 | |
| 104 | def test_v2_load_model_registers_moe_with_eplb(monkeypatch): |
| 105 | FakeEplbState.instances.clear() |
| 106 | model = SimpleNamespace(is_moe=True) |
| 107 | |
| 108 | monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler) |
| 109 | monkeypatch.setattr(eplb, "EplbState", FakeEplbState) |
| 110 | monkeypatch.setattr( |
| 111 | mrv2, |
| 112 | "get_model_loader", |
| 113 | lambda load_config: SimpleNamespace(load_model=lambda **_: model), |
| 114 | ) |
| 115 | monkeypatch.setattr( |
| 116 | mrv2, |
| 117 | "init_model_state", |
| 118 | lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), |
| 119 | ) |
| 120 | monkeypatch.setattr( |
| 121 | eplb, |
| 122 | "get_mixture_of_experts_model", |
| 123 | lambda loaded_model: ( |
| 124 | loaded_model if getattr(loaded_model, "is_moe", False) else None |
| 125 | ), |
| 126 | ) |
| 127 | |
| 128 | runner = _make_runner(is_last_pp_rank=False) |
| 129 | mrv2.GPUModelRunner.load_model(runner) |
| 130 | |
| 131 | assert runner.model is model |
| 132 | assert runner.model_state is not None |
| 133 | assert runner.eplb_state is not None |
| 134 | assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] |
| 135 | assert runner.eplb_state.async_started is True |
| 136 | |
| 137 | |
| 138 | def test_v2_load_model_with_dummy_weights_defers_eplb_async_loop(monkeypatch): |
| 139 | FakeEplbState.instances.clear() |
| 140 | model = SimpleNamespace(is_moe=True) |
| 141 | |
| 142 | monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler) |
| 143 | monkeypatch.setattr(eplb, "EplbState", FakeEplbState) |
| 144 | monkeypatch.setattr( |
| 145 | mrv2, |
| 146 | "get_model_loader", |
| 147 | lambda load_config: SimpleNamespace(load_model=lambda **_: model), |
| 148 | ) |
| 149 | monkeypatch.setattr( |
| 150 | mrv2, |
| 151 | "init_model_state", |
| 152 | lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), |
| 153 | ) |
| 154 | monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model) |
| 155 | |
| 156 | runner = _make_runner(is_last_pp_rank=False) |
| 157 | mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True) |
| 158 | |
| 159 | assert runner.load_config.load_format == "dummy" |
| 160 | assert runner.eplb_state is not None |
| 161 | # A scaling-up worker registers so it can join the EPLB communicator. |
| 162 | assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] |
| 163 | assert runner.eplb_state.async_started is False |
| 164 | |
| 165 | |
| 166 | def test_v2_setup_eplb_from_mapping_updates_state_in_place(): |
| 167 | runner = _make_runner() |
| 168 | state = FakeEplbState(runner.parallel_config, runner.device) |
| 169 | runner.eplb_state = state |
| 170 | mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64) |
| 171 | mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping) |
| 172 | |
| 173 | assert runner.eplb_state is state |
| 174 | assert state.update_mapping_calls == [(runner.model_config, mapping)] |
| 175 | |
| 176 | |
| 177 | def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): |
| 178 | events = [] |
| 179 | runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0) |
| 180 | runner.execute_model_state = SimpleNamespace( |
| 181 | input_batch=SimpleNamespace( |
| 182 | num_reqs=2, idx_mapping=torch.zeros(2, dtype=torch.int32) |
| 183 | ), |
| 184 | attn_metadata=None, |
| 185 | slot_mappings_by_layer=None, |
| 186 | hidden_states=None, |
| 187 | aux_hidden_states=None, |
| 188 | dp_sync=None, |
| 189 | finished_req_ids=set(), |
| 190 | ec_connector_output=None, |
| 191 | routed_experts=None, |
| 192 | cudagraph_stats=None, |
| 193 | ) |
| 194 | runner.req_states = SimpleNamespace() |
| 195 | |
| 196 | def fake_receive(*args, **kwargs): |
| 197 | events.append("receive") |
| 198 | # all_decode_next=True, so model_state.postprocess_state is skipped. |
| 199 | return True |
| 200 | |
| 201 | runner.pp_handler = SimpleNamespace(receive=fake_receive) |
| 202 | runner.postprocess_num_computed_tokens = lambda *args, **kwargs: events.append( |
| 203 | "postprocess_num_computed_tokens" |
| 204 | ) |
| 205 | runner.eplb.step = lambda *args, **kwargs: events.append("eplb") |
| 206 | |
| 207 | output = mrv2.GPUModelRunner.sample_tokens(runner, None) |
| 208 | assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) |
| 209 | assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] |