A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[V0 Deprecation] Remove V0 model runner base & simplify worker base (#25328)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Sep 20, 2025 at 20:49 UTC
1cd885bd540c3765e4c2d378b3167102dfa26ab5
4 files changed
+11
-706
vllm/attention/backends/abstract.py
+5
-10
@@ -4,19 +4,14 @@
4
from abc import ABC, abstractmethod
5
from contextlib import contextmanager
6
from dataclasses import dataclass, fields
7
-from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional,
8
- Protocol, Set, Tuple, Type, TypeVar)
7
+from typing import (Any, Dict, Generic, List, Optional, Protocol, Set, Tuple,
8
+ Type, TypeVar)
9
10
import torch
11
12
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
13
from vllm.multimodal import MultiModalPlaceholderMap
14
15
-if TYPE_CHECKING:
16
- from vllm.worker.model_runner_base import (ModelRunnerBase,
17
- ModelRunnerInputBase,
18
- ModelRunnerInputBuilderBase)
19
-
15
16
class AttentionType:
17
"""
@@ -170,7 +165,7 @@ class AttentionState(ABC, Generic[T]):
165
lifetime of the model runner."""
166
167
@abstractmethod
173
- def __init__(self, runner: "ModelRunnerBase"):
168
+ def __init__(self, runner: Any):
169
...
170
171
@abstractmethod
@@ -210,7 +205,7 @@ class AttentionState(ABC, Generic[T]):
205
...
206
207
@abstractmethod
213
- def begin_forward(self, model_input: "ModelRunnerInputBase") -> None:
208
+ def begin_forward(self, model_input) -> None:
209
"""Prepare state for forward pass."""
210
...
211
@@ -219,7 +214,7 @@ class AttentionMetadataBuilder(ABC, Generic[T]):
214
"""Abstract class for attention metadata builders."""
215
216
@abstractmethod
222
- def __init__(self, input_builder: "ModelRunnerInputBuilderBase") -> None:
217
+ def __init__(self, input_builder) -> None:
218
"""Create the builder, remember some configuration and parameters."""
219
raise NotImplementedError
220
vllm/attention/backends/utils.py
+2
-6
@@ -5,8 +5,7 @@ from collections import defaultdict
5
from contextlib import contextmanager
6
from dataclasses import dataclass
7
from itertools import accumulate
8
-from typing import (TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type,
9
- TypeVar, Union)
8
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
9
10
import numpy as np
11
import torch
@@ -21,9 +20,6 @@ from vllm.utils import async_tensor_h2d, make_tensor_with_pad
20
21
logger = init_logger(__name__)
22
24
-if TYPE_CHECKING:
25
- from vllm.worker.model_runner_base import ModelRunnerBase
26
-
23
# Error string(s) for encoder/decoder
24
# unsupported attention scenarios
25
STR_NOT_IMPL_ENC_DEC_ROCM_HIP = ("ROCm/HIP is not currently supported "
@@ -286,7 +282,7 @@ class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]):
282
283
class CommonAttentionState(AttentionState):
284
289
- def __init__(self, runner: "ModelRunnerBase"):
285
+ def __init__(self, runner):
286
self.runner = runner
287
self._is_graph_capturing = False
288
vllm/worker/model_runner_base.py
deleted
-307
@@ -1,307 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import dataclasses
5
-from abc import ABC, abstractmethod
6
-from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Type,
7
- TypeVar)
8
-
9
-import torch
10
-import torch.nn as nn
11
-
12
-from vllm.config import VllmConfig
13
-from vllm.logger import init_logger
14
-from vllm.model_executor.layers.sampler import SamplerOutput
15
-from vllm.model_executor.models.interfaces import supports_transcription
16
-from vllm.model_executor.models.interfaces_base import is_text_generation_model
17
-from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
18
-from vllm.tasks import GenerationTask, SupportedTask
19
-
20
-if TYPE_CHECKING:
21
- from vllm.attention import AttentionMetadata
22
- from vllm.attention.backends.abstract import AttentionBackend
23
- from vllm.model_executor import SamplingMetadata
24
-
25
-logger = init_logger(__name__)
26
-
27
-T = TypeVar('T', bound="BroadcastableModelInput")
28
-
29
-
30
-def _add_attn_metadata_broadcastable_dict(
31
- tensor_dict: Dict[str, Any],
32
- attn_metadata: Optional["AttentionMetadata"]) -> None:
33
- """
34
- Helper method to update tensor_dict with broadcastable
35
- AttentionMetadata fields.
36
- """
37
- if attn_metadata is not None:
38
- tensor_dict.update(attn_metadata.asdict_zerocopy())
39
-
40
-
41
-def _init_attn_metadata_from_tensor_dict(
42
- attn_backend: "AttentionBackend",
43
- tensor_dict: Dict[str, Any],
44
-) -> Dict[str, Any]:
45
- """
46
- Helper method to initialize AttentionMetadata based on an
47
- AttentionBackend and broadcastable AttentionMetadata fields.
48
- """
49
- # Extract the fields used to create AttentionMetadata.
50
- valid_attn_kwargs = {}
51
- for field in dataclasses.fields(attn_backend.get_metadata_cls()):
52
- if field.name in tensor_dict:
53
- if field.name == "input_positions":
54
- valid_attn_kwargs[field.name] = tensor_dict[field.name]
55
- else:
56
- valid_attn_kwargs[field.name] = tensor_dict.pop(field.name)
57
-
58
- attn_metadata = attn_backend.make_metadata(**valid_attn_kwargs)
59
- tensor_dict["attn_metadata"] = attn_metadata
60
- return tensor_dict
61
-
62
-
63
-def _init_sampling_metadata_from_tensor_dict( # type: ignore
64
- tensor_dict: Dict[str, Any]) -> Dict[str, Any]:
65
- """
66
- Helper method to initialize SamplingMetadata based on broadcastable
67
- SamplingMetadata fields.
68
- """
69
- from vllm.model_executor import SamplingMetadata
70
-
71
- selected_token_indices = tensor_dict.pop("selected_token_indices", None)
72
- # An empty SamplingMetadata to signal that the worker should skip
73
- # sampling.
74
- if selected_token_indices is not None:
75
- tensor_dict["sampling_metadata"] = SamplingMetadata(
76
- seq_groups=None,
77
- selected_token_indices=selected_token_indices,
78
- categorized_sample_indices=None,
79
- num_prompts=0,
80
- )
81
- return tensor_dict
82
-
83
-
84
-def _add_sampling_metadata_broadcastable_dict(
85
- tensor_dict: Dict[str, Any],
86
- sampling_metadata: Optional["SamplingMetadata"]) -> None:
87
- """
88
- Helper method to update tensor_dict with broadcastable
89
- SamplingMetadata fields.
90
- """
91
- if sampling_metadata is not None:
92
- tensor_dict["selected_token_indices"] = (
93
- sampling_metadata.selected_token_indices)
94
-
95
-
96
-def _init_frozen_model_input_from_tensor_dict(
97
- frozen_model_input_cls: Type["ModelRunnerInputBase"],
98
- tensor_dict: Dict[str, Any]) -> Dict[str, Any]:
99
- """
100
- Helper method to initialize a frozen ModelInput based on broadcastable
101
- """
102
- valid_tensor_kwargs = {}
103
- for field in dataclasses.fields(frozen_model_input_cls):
104
- val = tensor_dict.pop(field.name, None)
105
- if val is not None:
106
- valid_tensor_kwargs[field.name] = val
107
-
108
- frozen_model_input = frozen_model_input_cls(**valid_tensor_kwargs)
109
- tensor_dict["frozen_model_input"] = frozen_model_input
110
- return tensor_dict
111
-
112
-
113
-class BroadcastableModelInput(ABC):
114
-
115
- @abstractmethod
116
- def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
117
- """
118
- Extract broadcastable fields. Override for fields that require some
119
- custom deserialization.
120
- """
121
- raise NotImplementedError
122
-
123
- @classmethod
124
- @abstractmethod
125
- def from_broadcasted_tensor_dict(
126
- cls: Type[T],
127
- tensor_dict: Dict[str, Any],
128
- attn_backend: Optional["AttentionBackend"] = None,
129
- ) -> T:
130
- """
131
- Pop fields from the given tensor_dict and populate a new instance of
132
- BroadcastableModelInput.
133
- """
134
- raise NotImplementedError
135
-
136
-
137
-@dataclasses.dataclass(frozen=True)
138
-class ModelRunnerInputBase(BroadcastableModelInput):
139
- """Local inputs to each worker's model runner. May contain
140
- device-specific data. Different worker backends may have different methods
141
- of converting from the global ExecuteModelRequest produced by the LLM
142
- engine to the worker-local ModelRunnerInputBase objects.
143
-
144
- Model runners that support multi-GPU execution should define a
145
- ModelRunnerInputBase subclass, add their required fields, and specify how to
146
- serialize/deserialize a ModelInput for broadcast between workers.
147
- """
148
- pass
149
-
150
-
151
-class ModelRunnerInputBuilderBase(ABC, Generic[T]):
152
- """A builder to create ModelRunnerInputBase objects.
153
- """
154
-
155
- @abstractmethod
156
- def prepare(self,
157
- finished_requests_ids: Optional[List[str]] = None) -> None:
158
- raise NotImplementedError
159
-
160
- @abstractmethod
161
- def add_seq_group(self, seq_group_metadata):
162
- """TBA"""
163
- raise NotImplementedError
164
-
165
- @abstractmethod
166
- def build(self, *args, **kwargs) -> T:
167
- """Build metadata with on-device tensors."""
168
- raise NotImplementedError
169
-
170
-
171
-class ModelRunnerBase(ABC, Generic[T]):
172
- """
173
- Model runner interface that abstracts a particular hardware and/or type of
174
- model. Model execution may communicate data with model runners in other
175
- processes, but it should not include control plane metadata communication.
176
-
177
- Each ModelRunnerBase subclass should define a corresponding
178
- ModelRunnerInputBase subclass.
179
- """
180
-
181
- def __init__(
182
- self,
183
- vllm_config: VllmConfig,
184
- ) -> None:
185
- self.vllm_config = vllm_config
186
- self.model_config = vllm_config.model_config
187
- self.cache_config = vllm_config.cache_config
188
- self.lora_config = vllm_config.lora_config
189
- self.load_config = vllm_config.load_config
190
- self.parallel_config = vllm_config.parallel_config
191
- self.scheduler_config = vllm_config.scheduler_config
192
- self.device_config = vllm_config.device_config
193
- self.speculative_config = vllm_config.speculative_config
194
- self.observability_config = vllm_config.observability_config
195
-
196
- # Map of request_id -> generator used for seeded random sampling
197
- generators: Dict[str, torch.Generator] = {}
198
-
199
- @abstractmethod
200
- def make_model_input_from_broadcasted_tensor_dict(
201
- self,
202
- tensor_dict: Dict[str, Any],
203
- ) -> T:
204
- """
205
- Make an instance of a ModelRunnerInputBase from the broadcasted tensor
206
- dict.
207
- """
208
- raise NotImplementedError
209
-
210
- @abstractmethod
211
- def prepare_model_input(
212
- self,
213
- seq_group_metadata_list: List[SequenceGroupMetadata],
214
- virtual_engine: int = 0,
215
- finished_requests_ids: Optional[List[str]] = None,
216
- ) -> T:
217
- """
218
- Prepare the inputs to ModelRunnerBase.execute_model from an execution
219
- request. This method may move data to the worker's local device. It is
220
- not allowed to communicate with other workers or devices.
221
- """
222
- raise NotImplementedError
223
-
224
- @abstractmethod
225
- def get_model(self) -> nn.Module:
226
- raise NotImplementedError
227
-
228
- def get_supported_generation_tasks(self) -> list[GenerationTask]:
229
- model = self.get_model()
230
- supported_tasks = list[GenerationTask]()
231
-
232
- if is_text_generation_model(model):
233
- supported_tasks.append("generate")
234
-
235
- if supports_transcription(model):
236
- if model.supports_transcription_only:
237
- return ["transcription"]
238
-
239
- supported_tasks.append("transcription")
240
-
241
- return supported_tasks
242
-
243
- def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
244
- tasks = list[SupportedTask]()
245
-
246
- if self.model_config.runner_type == "generate":
247
- tasks.extend(self.get_supported_generation_tasks())
248
-
249
- return tuple(tasks)
250
-
251
- def execute_model(
252
- self,
253
- model_input: T,
254
- kv_caches: Optional[List[torch.Tensor]],
255
- intermediate_tensors: Optional[IntermediateTensors] = None,
256
- num_steps: int = 1,
257
- **kwargs,
258
- ) -> Optional[List[SamplerOutput]]:
259
- """
260
- Execute the model on the given input.
261
- """
262
- raise NotImplementedError
263
-
264
- def get_generators(self, finished_request_ids: Optional[List[str]] = None):
265
- """
266
- Return dict of per-request generators used for random sampling.
267
- """
268
-
269
- # Clean up generators from completed requests
270
- if finished_request_ids:
271
- for request_id in finished_request_ids:
272
- self.generators.pop(request_id, None)
273
-
274
- return self.generators
275
-
276
-
277
-class ModelRunnerWrapperBase:
278
- """
279
- The whole point of this class is to lazily initialize the model_runner.
280
- """
281
-
282
- def __init__(
283
- self,
284
- model_runner: ModelRunnerBase,
285
- ) -> None:
286
- self.model_runner: ModelRunnerBase = model_runner
287
-
288
- def __getattr__(self, attr):
289
- return getattr(self.model_runner, attr)
290
-
291
-
292
-class InputProcessingError(Exception):
293
- """This exception is raised when an error occurs preparing the inputs for
294
- a single sequence group.
295
- This allows the engine to gracefully handle errors with a single sequence
296
- group without having to fail the entire batch.
297
- """
298
-
299
- def __init__(self, request_id, message):
300
- """request_id is the id of the offending sequence group"""
301
- self.request_id = request_id
302
- self.message = message
303
- super().__init__(self.message)
304
-
305
- def __str__(self):
306
- return "Failed to prepare inputs for sequence group with request id: " \
307
- f"{self.request_id}, Error: {self.message}"
vllm/worker/worker_base.py
+4
-383
@@ -1,31 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import dataclasses
4
import os
6
-import time
7
-from abc import abstractmethod
8
-from typing import (Any, Callable, Dict, List, Optional, Set, Tuple, Type,
9
- TypeVar, Union)
5
+from typing import (Any, Callable, Dict, List, Optional, Set, Tuple, TypeVar,
6
+ Union)
7
8
import cloudpickle
12
-import torch
9
import torch.nn as nn
10
15
-from vllm.config import (ObservabilityConfig, VllmConfig,
16
- set_current_vllm_config)
17
-from vllm.distributed import broadcast_tensor_dict, get_pp_group, get_tp_group
11
+from vllm.config import VllmConfig, set_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.lora.request import LoRARequest
14
from vllm.model_executor.layers.sampler import SamplerOutput
21
-from vllm.sequence import ExecuteModelRequest, IntermediateTensors
15
+from vllm.sequence import ExecuteModelRequest
16
from vllm.utils import (enable_trace_function_call_for_thread,
17
resolve_obj_by_qualname, run_method,
18
update_environment_variables,
19
warn_for_unimplemented_methods)
26
-from vllm.worker.model_runner_base import (BroadcastableModelInput,
27
- ModelRunnerBase,
28
- ModelRunnerInputBase)
20
21
logger = init_logger(__name__)
22
@@ -141,356 +132,6 @@ class WorkerBase:
132
return
133
134
144
-class DelegateWorkerBase(WorkerBase):
145
- """
146
- A class that delegates all methods to another WorkerBase instance. This is
147
- useful for creating a WorkerBase that wraps another WorkerBase instance,
148
- e.g. speculative decoding.
149
- """
150
- worker: WorkerBase
151
-
152
- def __init__(
153
- self,
154
- *args,
155
- **kwargs,
156
- ) -> None:
157
- vllm_config: VllmConfig = kwargs.get("vllm_config")
158
- cls = resolve_obj_by_qualname(vllm_config.parallel_config.worker_cls)
159
- self.worker = cls(*args, **kwargs)
160
-
161
- def init_device(self) -> None:
162
- self.worker.init_device()
163
-
164
- def determine_num_available_blocks(self) -> Tuple[int, int]:
165
- return self.worker.determine_num_available_blocks()
166
-
167
- def initialize_cache(self, num_gpu_blocks: int,
168
- num_cpu_blocks: int) -> None:
169
- self.worker.initialize_cache(num_gpu_blocks, num_cpu_blocks)
170
-
171
- def load_model(self) -> None:
172
- """Load model onto target device."""
173
- self.worker.load_model()
174
-
175
- def get_model(self) -> nn.Module:
176
- return self.worker.get_model()
177
-
178
- def execute_model(
179
- self,
180
- execute_model_req: Optional[ExecuteModelRequest] = None
181
- ) -> Optional[List[SamplerOutput]]:
182
- return self.worker.execute_model(execute_model_req)
183
-
184
- def get_cache_block_size_bytes(self) -> int:
185
- return self.worker.get_cache_block_size_bytes()
186
-
187
- def add_lora(self, lora_request: LoRARequest) -> bool:
188
- return self.worker.add_lora(lora_request)
189
-
190
- def remove_lora(self, lora_id: int) -> bool:
191
- return self.worker.remove_lora(lora_id)
192
-
193
- def pin_lora(self, lora_id: int) -> bool:
194
- return self.worker.pin_lora(lora_id)
195
-
196
- def list_loras(self) -> Set[int]:
197
- return self.worker.list_loras()
198
-
199
- def __getattr__(self, attr):
200
- return getattr(self.worker, attr)
201
-
202
-
203
-class LoRANotSupportedWorkerBase(WorkerBase):
204
- """Partial implementation of WorkerBase that raises exceptions when LoRA
205
- methods are invoked.
206
- """
207
-
208
- def add_lora(self, lora_request: LoRARequest) -> bool:
209
- raise ValueError(f"{type(self)} does not support LoRA")
210
-
211
- def remove_lora(self, lora_id: int) -> bool:
212
- raise ValueError(f"{type(self)} does not support LoRA")
213
-
214
- def pin_lora(self, lora_id: int) -> bool:
215
- raise ValueError(f"{type(self)} does not support LoRA")
216
-
217
- def list_loras(self) -> Set[int]:
218
- raise ValueError(f"{type(self)} does not support LoRA")
219
-
220
-
221
-@dataclasses.dataclass(frozen=True)
222
-class WorkerInput:
223
- """Local inputs to each worker. May contain device-specific data. These
224
- fields should be broadcastable to other workers.
225
- """
226
-
227
- num_seq_groups: Optional[int] = None
228
- blocks_to_swap_in: Optional[torch.Tensor] = None
229
- blocks_to_swap_out: Optional[torch.Tensor] = None
230
- blocks_to_copy: Optional[torch.Tensor] = None
231
- virtual_engine: int = 0
232
- num_steps: int = 1
233
-
234
- @classmethod
235
- def from_broadcasted_tensor_dict(
236
- cls: Type["WorkerInput"],
237
- tensor_dict: Dict[str, Any],
238
- ) -> "WorkerInput":
239
- """
240
- Pop fields from the given tensor_dict and populate a new instance of
241
- WorkerInput.
242
- """
243
- return cls(
244
- num_seq_groups=tensor_dict.pop("num_seq_groups"),
245
- blocks_to_swap_in=tensor_dict.pop("blocks_to_swap_in"),
246
- blocks_to_swap_out=tensor_dict.pop("blocks_to_swap_out"),
247
- blocks_to_copy=tensor_dict.pop("blocks_to_copy"),
248
- virtual_engine=tensor_dict["virtual_engine"],
249
- num_steps=tensor_dict.pop("num_steps"),
250
- )
251
-
252
- def as_broadcastable_tensor_dict(
253
- self) -> Dict[str, Union[int, torch.Tensor]]:
254
- """
255
- Extract broadcastable fields.
256
- """
257
- tensor_dict = {
258
- "num_seq_groups": self.num_seq_groups,
259
- "blocks_to_swap_in": self.blocks_to_swap_in,
260
- "blocks_to_swap_out": self.blocks_to_swap_out,
261
- "blocks_to_copy": self.blocks_to_copy,
262
- "virtual_engine": self.virtual_engine,
263
- "num_steps": self.num_steps,
264
- }
265
-
266
- return tensor_dict
267
-
268
-
269
-class LocalOrDistributedWorkerBase(WorkerBase):
270
- """
271
- Partial implementation of WorkerBase that has a default `execute_model`
272
- definition to perform metadata transfer between workers when in distributed
273
- mode. Subclasses of this interface should use model runners that inherit
274
- from ModelRunnerBase, and should only need to implement worker-local logic.
275
- If custom control plane logic is needed to transfer metadata, or if the
276
- model runner cannot inherit from ModelRunnerBase, use WorkerBase instead.
277
- """
278
- is_driver_worker: bool
279
- model_runner: ModelRunnerBase
280
- observability_config: Optional[ObservabilityConfig] = None
281
-
282
- @property
283
- @abstractmethod
284
- def do_metadata_broadcast(self) -> bool:
285
- """
286
- Used by the default `execute_model` to check whether broadcast is
287
- needed to transfer request inputs from the driver worker to other
288
- workers in the TP group. If WorkerBase subclass only supports
289
- single-worker execution, then this method should return False.
290
- """
291
- raise NotImplementedError
292
-
293
- @property
294
- @abstractmethod
295
- def kv_cache(self) -> Optional[List[List[torch.Tensor]]]:
296
- """
297
- Gets the list of kv caches to pass to the worker's model runner. Each
298
- element in the list is a kv cache corresponding to a particular virtual
299
- engine (PP stream). Used by the default `execute_model`. If the worker's
300
- model runner does not follow the ModelRunnerBase interface, then inherit
301
- from WorkerBase instead.
302
- """
303
- raise NotImplementedError
304
-
305
- @abstractmethod
306
- def prepare_worker_input(
307
- self, execute_model_req: ExecuteModelRequest) -> WorkerInput:
308
- """
309
- Prepare the inputs to WorkerBase.execute_worker from an execution
310
- request. This method may move data to the worker's local device. It is
311
- not allowed to communicate with other workers or devices.
312
- """
313
- raise NotImplementedError
314
-
315
- @abstractmethod
316
- def execute_worker(self, worker_input: WorkerInput) -> None:
317
- """
318
- Process an execution request.
319
- """
320
- raise NotImplementedError
321
-
322
- def _get_worker_input_from_broadcast(
323
- self
324
- ) -> Optional[Tuple[BroadcastableModelInput, WorkerInput, Dict[
325
- str, torch.Tensor]]]:
326
- """ Get the worker input from the broadcasted tensor dict. """
327
- assert self.do_metadata_broadcast
328
- assert not self.is_driver_worker
329
- broadcast_data = broadcast_tensor_dict(src=0)
330
- if not broadcast_data:
331
- return None
332
-
333
- worker_input = WorkerInput.from_broadcasted_tensor_dict(broadcast_data)
334
- model_input = (
335
- self.model_runner.make_model_input_from_broadcasted_tensor_dict(
336
- broadcast_data))
337
-
338
- kwargs = extract_previous_hidden_states(broadcast_data)
339
-
340
- return model_input, worker_input, kwargs
341
-
342
- def _get_driver_input_and_broadcast(
343
- self, execute_model_req: ExecuteModelRequest
344
- ) -> Tuple[BroadcastableModelInput, WorkerInput, Dict[str, torch.Tensor]]:
345
- """ Get the driver input and broadcast it to other workers. """
346
- assert self.is_driver_worker
347
-
348
- worker_input: WorkerInput = self.prepare_worker_input(
349
- execute_model_req=execute_model_req)
350
- model_input: ModelRunnerInputBase = (
351
- self.model_runner.prepare_model_input(
352
- execute_model_req.seq_group_metadata_list,
353
- execute_model_req.virtual_engine,
354
- execute_model_req.finished_requests_ids))
355
-
356
- kwargs = extract_previous_hidden_states(execute_model_req)
357
-
358
- if self.do_metadata_broadcast:
359
- broadcast_data = worker_input.as_broadcastable_tensor_dict()
360
- broadcast_data.update(model_input.as_broadcastable_tensor_dict())
361
- broadcast_data.update(kwargs)
362
- broadcast_tensor_dict(broadcast_data, src=0)
363
-
364
- if execute_model_req.async_callback:
365
- model_input = dataclasses.replace( # type: ignore
366
- model_input,
367
- async_callback=execute_model_req.async_callback)
368
-
369
- return model_input, worker_input, kwargs
370
-
371
- def prepare_input(
372
- self,
373
- execute_model_req: Optional[ExecuteModelRequest] = None
374
- ) -> Optional[Tuple[BroadcastableModelInput, WorkerInput, Dict[
375
- str, torch.Tensor]]]:
376
- """
377
- Prepare the inputs to ModelRunner and workers.
378
- """
379
- if self.is_driver_worker:
380
- if execute_model_req is None:
381
- if self.do_metadata_broadcast:
382
- # This signals that there's no more requests to process for
383
- # now. All workers are running infinite loop with
384
- # broadcast_tensor_dict, and it stops the loop when the
385
- # driver broadcasts an empty input. Send an empty input to
386
- # notify all other workers to stop their execution loop.
387
- broadcast_tensor_dict({}, src=0)
388
- return None
389
- return self._get_driver_input_and_broadcast(execute_model_req)
390
- else:
391
- return self._get_worker_input_from_broadcast()
392
-
393
- def get_model(self) -> nn.Module:
394
- return self.model_runner.get_model()
395
-
396
- def execute_model(
397
- self,
398
- execute_model_req: Optional[ExecuteModelRequest] = None,
399
- ) -> Optional[List[SamplerOutput]]:
400
- """Executes at least one model step on the given sequences, unless no
401
- sequences are provided."""
402
- start_time = time.perf_counter()
403
-
404
- inputs = self.prepare_input(execute_model_req)
405
- if inputs is None:
406
- return None
407
-
408
- model_input, worker_input, kwargs = inputs
409
- num_steps = worker_input.num_steps
410
-
411
- self.execute_worker(worker_input)
412
-
413
- # If there is no input, we don't need to execute the model.
414
- if worker_input.num_seq_groups == 0:
415
- return []
416
-
417
- intermediate_tensors = None
418
- orig_model_execute_time = 0.0
419
- if not get_pp_group().is_first_rank:
420
- intermediate_tensors = IntermediateTensors(
421
- get_pp_group().recv_tensor_dict(
422
- all_gather_group=get_tp_group()))
423
- if (self.observability_config is not None
424
- and self.observability_config.collect_model_execute_time):
425
- orig_model_execute_time = intermediate_tensors.tensors.get(
426
- "model_execute_time", torch.tensor(0)).item()
427
-
428
- output = self.model_runner.execute_model(
429
- model_input=model_input,
430
- kv_caches=self.kv_cache[worker_input.virtual_engine]
431
- if self.kv_cache is not None else None,
432
- intermediate_tensors=intermediate_tensors,
433
- num_steps=num_steps,
434
- **kwargs,
435
- )
436
-
437
- model_execute_time = time.perf_counter() - start_time
438
- if not get_pp_group().is_last_rank:
439
- # output is IntermediateTensors
440
- assert isinstance(output, IntermediateTensors)
441
- if (self.observability_config is not None
442
- and self.observability_config.collect_model_execute_time):
443
- output.tensors["model_execute_time"] = torch.tensor(
444
- model_execute_time + orig_model_execute_time)
445
- get_pp_group().send_tensor_dict(output.tensors,
446
- all_gather_group=get_tp_group())
447
- return [None]
448
- if (self.observability_config is not None
449
- and self.observability_config.collect_model_execute_time
450
- and output is not None):
451
- for o in output:
452
- o.model_execute_time = (orig_model_execute_time +
453
- model_execute_time)
454
-
455
- # output is List[SamplerOutput]
456
- return output
457
-
458
- def _execute_model_spmd(
459
- self,
460
- execute_model_req: ExecuteModelRequest,
461
- intermediate_tensors: Optional[IntermediateTensors] = None
462
- ) -> Optional[List[SamplerOutput]]:
463
- """
464
- Execute model in Single Program Multiple Data (SPMD) fashion.
465
- All workers take the same request, prepare the input and
466
- execute the model.
467
- """
468
- assert execute_model_req is not None, (
469
- "_execute_model_spmd() requires each worker to take in an "
470
- "ExecuteModelRequest")
471
- worker_input: WorkerInput = self.prepare_worker_input(
472
- execute_model_req=execute_model_req)
473
- model_input: ModelRunnerInputBase = (
474
- self.model_runner.prepare_model_input(
475
- execute_model_req.seq_group_metadata_list))
476
-
477
- self.execute_worker(worker_input)
478
-
479
- # If there is no input, we don't need to execute the model.
480
- if worker_input.num_seq_groups == 0:
481
- return []
482
-
483
- kwargs = extract_previous_hidden_states(execute_model_req)
484
-
485
- return self.model_runner.execute_model(
486
- model_input=model_input,
487
- kv_caches=self.kv_cache[worker_input.virtual_engine]
488
- if self.kv_cache is not None else None,
489
- intermediate_tensors=intermediate_tensors,
490
- **kwargs,
491
- )
492
-
493
-
135
class WorkerWrapperBase:
136
"""
137
This class represents one process in an executor/engine. It is responsible
@@ -636,23 +277,3 @@ class WorkerWrapperBase:
277
278
def __getattr__(self, attr):
279
return getattr(self.worker, attr)
639
-
640
-
641
-def extract_previous_hidden_states(
642
- data: Union[ExecuteModelRequest, Dict[str, torch.Tensor]]) -> \
643
- Dict[str, torch.Tensor]:
644
- """If data contains previous_hidden_states, extract it. This returns a dict
645
- which can be used directly as additional kwargs in any following
646
- execute_model calls. This is used in draft models like EAGLE."""
647
- output = {}
648
-
649
- # When called from non-driver worker, data is dict but when called from
650
- # driver worker, data is ExecuteModelRequest.
651
- if isinstance(data, dict):
652
- if "previous_hidden_states" in data:
653
- output["previous_hidden_states"] = data["previous_hidden_states"]
654
- elif data.previous_hidden_states is not None:
655
- output["previous_hidden_states"] = data.previous_hidden_states\
656
- .hidden_states
657
-
658
- return output