"tests/vscode:/vscode.git/clone" did not exist on "326976291b541f0fd5bef34aa1ff4a84bf8fb37d"
llm_engine.py 37 KB
Newer Older
Antoni Baum's avatar
Antoni Baum committed
1
import time
2
from typing import Iterable, List, Optional, Type, Union
3

4
from transformers import GenerationConfig, PreTrainedTokenizer
5

6
import vllm
7
8
9
from vllm.config import (CacheConfig, DecodingConfig, DeviceConfig, LoadConfig,
                         LoRAConfig, ModelConfig, ParallelConfig,
                         SchedulerConfig, SpeculativeConfig,
10
                         VisionLanguageConfig)
11
12
from vllm.core.scheduler import (ScheduledSequenceGroup, Scheduler,
                                 SchedulerOutputs)
Woosuk Kwon's avatar
Woosuk Kwon committed
13
from vllm.engine.arg_utils import EngineArgs
14
from vllm.engine.metrics import StatLogger, Stats
15
16
17
18
from vllm.engine.output_processor.interfaces import (
    SequenceGroupOutputProcessor)
from vllm.engine.output_processor.stop_checker import StopChecker
from vllm.engine.output_processor.util import create_output_by_sequence_group
19
from vllm.executor.executor_base import ExecutorBase
20
from vllm.executor.ray_utils import initialize_ray_cluster
Woosuk Kwon's avatar
Woosuk Kwon committed
21
from vllm.logger import init_logger
22
from vllm.lora.request import LoRARequest
23
24
25
from vllm.outputs import (EmbeddingRequestOutput, RequestOutput,
                          RequestOutputFactory)
from vllm.pooling_params import PoolingParams
Woosuk Kwon's avatar
Woosuk Kwon committed
26
from vllm.sampling_params import SamplingParams
27
28
from vllm.sequence import (EmbeddingSequenceGroupOutput, ExecuteModelRequest,
                           MultiModalData, PoolerOutput, SamplerOutput,
29
                           Sequence, SequenceGroup, SequenceGroupMetadata,
30
                           SequenceStatus)
31
from vllm.transformers_utils.detokenizer import Detokenizer
32
33
from vllm.transformers_utils.tokenizer_group import (BaseTokenizerGroup,
                                                     get_tokenizer_group)
yhu422's avatar
yhu422 committed
34
35
from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,
                                  usage_message)
36
from vllm.utils import Counter
37
38

logger = init_logger(__name__)
39
_LOCAL_LOGGING_INTERVAL_SEC = 5
Woosuk Kwon's avatar
Woosuk Kwon committed
40

41

42
43
44
45
46
47
48
49
50
51
52
def _load_generation_config_dict(model_config: ModelConfig):
    try:
        return GenerationConfig.from_pretrained(
            model_config.model,
            revision=model_config.revision,
        ).to_diff_dict()
    except OSError:
        # Not found.
        return {}


53
class LLMEngine:
Zhuohan Li's avatar
Zhuohan Li committed
54
    """An LLM engine that receives requests and generates texts.
55

Woosuk Kwon's avatar
Woosuk Kwon committed
56
    This is the main class for the vLLM engine. It receives requests
57
58
59
60
61
62
63
    from clients and generates texts from the LLM. It includes a tokenizer, a
    language model (possibly distributed across multiple GPUs), and GPU memory
    space allocated for intermediate states (aka KV cache). This class utilizes
    iteration-level scheduling and efficient memory management to maximize the
    serving throughput.

    The `LLM` class wraps this class for offline batched inference and the
64
    `AsyncLLMEngine` class wraps this class for online serving.
65

Zhuohan Li's avatar
Zhuohan Li committed
66
67
    NOTE: The config arguments are derived from the `EngineArgs` class. For the
    comprehensive list of arguments, see `EngineArgs`.
68
69
70
71
72
73
74

    Args:
        model_config: The configuration related to the LLM model.
        cache_config: The configuration related to the KV cache memory
            management.
        parallel_config: The configuration related to distributed execution.
        scheduler_config: The configuration related to the request scheduler.
75
        device_config: The configuration related to the device.
76
77
78
79
80
        lora_config (Optional): The configuration related to serving multi-LoRA.
        vision_language_config (Optional): The configuration related to vision
            language models.
        speculative_config (Optional): The configuration related to speculative
            decoding.
81
82
        executor_class: The model executor class for managing distributed
            execution.
83
        log_stats: Whether to log statistics.
yhu422's avatar
yhu422 committed
84
        usage_context: Specified entry point, used for usage info collection
85
    """
86
87
88
89
90
91
92

    def __init__(
        self,
        model_config: ModelConfig,
        cache_config: CacheConfig,
        parallel_config: ParallelConfig,
        scheduler_config: SchedulerConfig,
93
        device_config: DeviceConfig,
94
        load_config: LoadConfig,
95
        lora_config: Optional[LoRAConfig],
96
97
        vision_language_config: Optional[VisionLanguageConfig],
        speculative_config: Optional[SpeculativeConfig],
98
        decoding_config: Optional[DecodingConfig],
99
        executor_class: Type[ExecutorBase],
100
        log_stats: bool,
yhu422's avatar
yhu422 committed
101
        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
102
103
    ) -> None:
        logger.info(
104
105
106
            "Initializing an LLM engine (v%s) with config: "
            "model=%r, speculative_config=%r, tokenizer=%r, "
            "skip_tokenizer_init=%s, tokenizer_mode=%s, revision=%s, "
107
108
109
110
111
            "rope_scaling=%r, tokenizer_revision=%s, "
            "trust_remote_code=%s, dtype=%s, max_seq_len=%d, "
            "download_dir=%r, load_format=%s, tensor_parallel_size=%d, "
            "disable_custom_all_reduce=%s, quantization=%s, "
            "enforce_eager=%s, kv_cache_dtype=%s, "
112
            "quantization_param_path=%s, device_config=%s, "
113
            "decoding_config=%r, seed=%d, served_model_name=%s)",
114
115
116
117
118
119
120
            vllm.__version__,
            model_config.model,
            speculative_config,
            model_config.tokenizer,
            model_config.skip_tokenizer_init,
            model_config.tokenizer_mode,
            model_config.revision,
121
            model_config.rope_scaling,
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
            model_config.tokenizer_revision,
            model_config.trust_remote_code,
            model_config.dtype,
            model_config.max_model_len,
            load_config.download_dir,
            load_config.load_format,
            parallel_config.tensor_parallel_size,
            parallel_config.disable_custom_all_reduce,
            model_config.quantization,
            model_config.enforce_eager,
            cache_config.cache_dtype,
            model_config.quantization_param_path,
            device_config.device,
            decoding_config,
            model_config.seed,
137
            model_config.served_model_name,
138
        )
139
140
141
142
        # TODO(woosuk): Print more configs in debug mode.

        self.model_config = model_config
        self.cache_config = cache_config
143
        self.lora_config = lora_config
144
        self.vision_language_config = vision_language_config
145
146
        self.parallel_config = parallel_config
        self.scheduler_config = scheduler_config
147
        self.device_config = device_config
148
        self.speculative_config = speculative_config
149
        self.load_config = load_config
150
        self.decoding_config = decoding_config or DecodingConfig()
151
152
        self.log_stats = log_stats

153
154
155
156
157
158
159
160
        if not self.model_config.skip_tokenizer_init:
            self.tokenizer: BaseTokenizerGroup
            self._init_tokenizer()
            self.detokenizer = Detokenizer(self.tokenizer)
        else:
            self.detokenizer = None
            self.tokenizer = None

161
        self.seq_counter = Counter()
162
163
        self.generation_config_fields = _load_generation_config_dict(
            model_config)
164

165
166
167
168
169
170
171
172
173
        self.model_executor = executor_class(
            model_config=model_config,
            cache_config=cache_config,
            parallel_config=parallel_config,
            scheduler_config=scheduler_config,
            device_config=device_config,
            lora_config=lora_config,
            vision_language_config=vision_language_config,
            speculative_config=speculative_config,
174
            load_config=load_config,
175
        )
176

177
178
        if not self.model_config.embedding_mode:
            self._initialize_kv_caches()
179

yhu422's avatar
yhu422 committed
180
181
        # If usage stat is enabled, collect relevant info.
        if is_usage_stats_enabled():
182
183
            from vllm.model_executor.model_loader import (
                get_architecture_class_name)
yhu422's avatar
yhu422 committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
            usage_message.report_usage(
                get_architecture_class_name(model_config),
                usage_context,
                extra_kvs={
                    # Common configuration
                    "dtype":
                    str(model_config.dtype),
                    "tensor_parallel_size":
                    parallel_config.tensor_parallel_size,
                    "block_size":
                    cache_config.block_size,
                    "gpu_memory_utilization":
                    cache_config.gpu_memory_utilization,

                    # Quantization
                    "quantization":
                    model_config.quantization,
                    "kv_cache_dtype":
                    cache_config.cache_dtype,

                    # Feature flags
                    "enable_lora":
                    bool(lora_config),
                    "enable_prefix_caching":
                    cache_config.enable_prefix_caching,
                    "enforce_eager":
                    model_config.enforce_eager,
                    "disable_custom_all_reduce":
                    parallel_config.disable_custom_all_reduce,
                })

215
216
217
218
        if self.tokenizer:
            # Ping the tokenizer to ensure liveness if it runs in a
            # different process.
            self.tokenizer.ping()
219

220
        # Create the scheduler.
221
222
        # NOTE: the cache_config here have been updated with the numbers of
        # GPU and CPU blocks, which are profiled in the distributed executor.
223
        self.scheduler = Scheduler(scheduler_config, cache_config, lora_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
224

225
226
227
        # Metric Logging.
        if self.log_stats:
            self.stat_logger = StatLogger(
228
                local_interval=_LOCAL_LOGGING_INTERVAL_SEC,
229
                labels=dict(model_name=model_config.served_model_name),
230
                max_model_len=self.model_config.max_model_len)
231
            self.stat_logger.info("cache_config", self.cache_config)
232

233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
        # Create sequence output processor, e.g. for beam search or
        # speculative decoding.
        self.output_processor = (
            SequenceGroupOutputProcessor.create_output_processor(
                self.scheduler_config,
                self.detokenizer,
                self.scheduler,
                self.seq_counter,
                self.get_tokenizer_for_seq,
                stop_checker=StopChecker(
                    self.scheduler_config.max_model_len,
                    self.get_tokenizer_for_seq,
                ),
            ))

248
249
250
251
252
253
254
255
256
257
258
    def _initialize_kv_caches(self) -> None:
        """Initialize the KV cache in the worker(s).

        The workers will determine the number of blocks in both the GPU cache
        and the swap CPU cache.
        """
        num_gpu_blocks, num_cpu_blocks = (
            self.model_executor.determine_num_available_blocks())

        if self.cache_config.num_gpu_blocks_override is not None:
            num_gpu_blocks_override = self.cache_config.num_gpu_blocks_override
259
260
261
262
            logger.info(
                "Overriding num_gpu_blocks=%d with "
                "num_gpu_blocks_override=%d", num_gpu_blocks,
                num_gpu_blocks_override)
263
264
265
266
267
268
269
            num_gpu_blocks = num_gpu_blocks_override

        self.cache_config.num_gpu_blocks = num_gpu_blocks
        self.cache_config.num_cpu_blocks = num_cpu_blocks

        self.model_executor.initialize_cache(num_gpu_blocks, num_cpu_blocks)

270
    @classmethod
yhu422's avatar
yhu422 committed
271
272
273
274
275
    def from_engine_args(
        cls,
        engine_args: EngineArgs,
        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
    ) -> "LLMEngine":
276
277
        """Creates an LLM engine from the engine arguments."""
        # Create the engine configs.
278
        engine_config = engine_args.create_engine_config()
279
280
        distributed_executor_backend = (
            engine_config.parallel_config.distributed_executor_backend)
281
282

        # Initialize the cluster and specify the executor class.
283
        if engine_config.device_config.device_type == "neuron":
284
285
            from vllm.executor.neuron_executor import NeuronExecutor
            executor_class = NeuronExecutor
286
        elif engine_config.device_config.device_type == "cpu":
287
288
            from vllm.executor.cpu_executor import CPUExecutor
            executor_class = CPUExecutor
289
        elif distributed_executor_backend == "ray":
290
            initialize_ray_cluster(engine_config.parallel_config)
291
292
            from vllm.executor.ray_gpu_executor import RayGPUExecutor
            executor_class = RayGPUExecutor
293
294
295
296
        elif distributed_executor_backend == "mp":
            from vllm.executor.multiproc_gpu_executor import (
                MultiprocessingGPUExecutor)
            executor_class = MultiprocessingGPUExecutor
297
298
299
300
301
        else:
            from vllm.executor.gpu_executor import GPUExecutor
            executor_class = GPUExecutor

        # Create the LLM engine.
yhu422's avatar
yhu422 committed
302
        engine = cls(
303
            **engine_config.to_dict(),
yhu422's avatar
yhu422 committed
304
305
306
307
            executor_class=executor_class,
            log_stats=not engine_args.disable_log_stats,
            usage_context=usage_context,
        )
308
        return engine
309

310
311
312
313
314
    def __reduce__(self):
        # This is to ensure that the LLMEngine is not referenced in
        # the closure used to initialize Ray worker actors
        raise RuntimeError("LLMEngine should not be pickled!")

315
316
317
318
319
320
    def __del__(self):
        # Shutdown model executor when engine is garbage collected
        # Use getattr since __init__ can fail before the field is set
        if model_executor := getattr(self, "model_executor", None):
            model_executor.shutdown()

321
    def get_tokenizer(self) -> "PreTrainedTokenizer":
322
        return self.tokenizer.get_lora_tokenizer(None)
323
324
325

    def get_tokenizer_for_seq(self,
                              sequence: Sequence) -> "PreTrainedTokenizer":
326
327
328
329
        return self.tokenizer.get_lora_tokenizer(sequence.lora_request)

    def _init_tokenizer(self, **tokenizer_init_kwargs):
        init_kwargs = dict(
330
            tokenizer_id=self.model_config.tokenizer,
331
332
333
334
335
336
337
            enable_lora=bool(self.lora_config),
            max_num_seqs=self.scheduler_config.max_num_seqs,
            max_input_length=None,
            tokenizer_mode=self.model_config.tokenizer_mode,
            trust_remote_code=self.model_config.trust_remote_code,
            revision=self.model_config.tokenizer_revision)
        init_kwargs.update(tokenizer_init_kwargs)
338
        self.tokenizer = get_tokenizer_group(
339
            self.parallel_config.tokenizer_pool_config, **init_kwargs)
340

341
342
    def _verify_args(self) -> None:
        self.model_config.verify_with_parallel_config(self.parallel_config)
343
        self.cache_config.verify_with_parallel_config(self.parallel_config)
344
345
346
347
        if self.lora_config:
            self.lora_config.verify_with_model_config(self.model_config)
            self.lora_config.verify_with_scheduler_config(
                self.scheduler_config)
348

349
350
351
352
353
354
355
356
357
358
359
360
361
362
    def encode_request(
        self,
        request_id: str,  # pylint: disable=unused-argument
        prompt: Optional[str],
        prompt_token_ids: Optional[List[int]] = None,
        lora_request: Optional[LoRARequest] = None,
    ):
        if prompt_token_ids is None:
            assert prompt is not None
            prompt_token_ids = self.tokenizer.encode(request_id=request_id,
                                                     prompt=prompt,
                                                     lora_request=lora_request)
        return prompt_token_ids

363
364
365
    def add_request(
        self,
        request_id: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
366
        prompt: Optional[str],
367
        params: Union[SamplingParams, PoolingParams],
368
369
        prompt_token_ids: Optional[List[int]] = None,
        arrival_time: Optional[float] = None,
370
        lora_request: Optional[LoRARequest] = None,
371
        multi_modal_data: Optional[MultiModalData] = None,
372
    ) -> None:
Zhuohan Li's avatar
Zhuohan Li committed
373
        """Add a request to the engine's request pool.
374
375

        The request is added to the request pool and will be processed by the
Zhuohan Li's avatar
Zhuohan Li committed
376
        scheduler as `engine.step()` is called. The exact scheduling policy is
377
378
379
380
381
382
        determined by the scheduler.

        Args:
            request_id: The unique ID of the request.
            prompt: The prompt string. Can be None if prompt_token_ids is
                provided.
383
384
            params: Parameters for sampling or pooling. SamplingParams
                for text generation. PoolingParams for pooling.
385
386
387
            prompt_token_ids: The token IDs of the prompt. If None, we
                use the tokenizer to convert the prompts to token IDs.
            arrival_time: The arrival time of the request. If None, we use
388
                the current monotonic time.
389
            multi_modal_data: Multi modal data per request.
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413

        Details:
            - Set arrival_time to the current time if it is None.
            - Set prompt_token_ids to the encoded prompt if it is None.
            - Create `best_of` number of :class:`~vllm.Sequence` objects.
            - Create a :class:`~vllm.SequenceGroup` object
              from the list of :class:`~vllm.Sequence`.
            - Add the :class:`~vllm.SequenceGroup` object to the scheduler.

        Example:
            >>> # initialize engine
            >>> engine = LLMEngine.from_engine_args(engine_args)
            >>> # set request arguments
            >>> example_prompt = "Who is the president of the United States?"
            >>> sampling_params = SamplingParams(temperature=0.0)
            >>> request_id = 0
            >>>
            >>> # add the request to the engine
            >>> engine.add_request(
            >>>    str(request_id),
            >>>    example_prompt,
            >>>    SamplingParams(temperature=0.0))
            >>> # continue the request processing
            >>> ...
414
        """
415
416
417
        if lora_request is not None and not self.lora_config:
            raise ValueError(f"Got lora_request {lora_request} but LoRA is "
                             "not enabled!")
418
        if arrival_time is None:
419
            arrival_time = time.time()
420
421
422
423
424
        prompt_token_ids = self.encode_request(
            request_id=request_id,
            prompt=prompt,
            prompt_token_ids=prompt_token_ids,
            lora_request=lora_request)
425
426
427

        # Create the sequences.
        block_size = self.cache_config.block_size
428
        seq_id = next(self.seq_counter)
429
430
431
432
433
434
435
        eos_token_id = None
        if self.tokenizer:
            eos_token_id = self.tokenizer.get_lora_tokenizer(
                lora_request).eos_token_id
        else:
            logger.warning("Use None for EOS token id because tokenizer is "
                           "not initialized")
436
        seq = Sequence(seq_id, prompt, prompt_token_ids, block_size,
437
                       eos_token_id, lora_request)
438

439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
        # Create a SequenceGroup based on SamplingParams or PoolingParams
        if isinstance(params, SamplingParams):
            seq_group = self._create_sequence_group_with_sampling(
                request_id,
                seq,
                params,
                arrival_time,
                lora_request,
                multi_modal_data,
            )
        elif isinstance(params, PoolingParams):
            seq_group = self._create_sequence_group_with_pooling(
                request_id,
                seq,
                params,
                arrival_time,
                lora_request,
                multi_modal_data,
            )
        else:
            raise ValueError(
                "Either SamplingParams or PoolingParams must be provided.")

        # Add the sequence group to the scheduler.
        self.scheduler.add_seq_group(seq_group)

    def _create_sequence_group_with_sampling(
        self,
        request_id: str,
        seq: Sequence,
        sampling_params: SamplingParams,
        arrival_time: Optional[float] = None,
        lora_request: Optional[LoRARequest] = None,
        multi_modal_data: Optional[MultiModalData] = None,
    ) -> SequenceGroup:
        """Creates a SequenceGroup with SamplingParams."""
        max_logprobs = self.get_model_config().max_logprobs
        if (sampling_params.logprobs
                and sampling_params.logprobs > max_logprobs) or (
                    sampling_params.prompt_logprobs
                    and sampling_params.prompt_logprobs > max_logprobs):
            raise ValueError(f"Cannot request more than "
                             f"{max_logprobs} logprobs.")

483
484
485
        # Defensive copy of SamplingParams, which are used by the sampler,
        # this doesn't deep-copy LogitsProcessor objects
        sampling_params = sampling_params.clone()
486
        # Add the eos token id into the sampling_params to support min_tokens
487
        # processing
488
489
        if seq.eos_token_id is not None:
            sampling_params.all_stop_token_ids.add(seq.eos_token_id)
490
491
        sampling_params.update_from_generation_config(
            self.generation_config_fields)
492

493
        # Create the sequence group.
494
495
496
497
498
499
        seq_group = SequenceGroup(request_id=request_id,
                                  seqs=[seq],
                                  arrival_time=arrival_time,
                                  sampling_params=sampling_params,
                                  lora_request=lora_request,
                                  multi_modal_data=multi_modal_data)
500

501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
        return seq_group

    def _create_sequence_group_with_pooling(
        self,
        request_id: str,
        seq: Sequence,
        pooling_params: PoolingParams,
        arrival_time: Optional[float] = None,
        lora_request: Optional[LoRARequest] = None,
        multi_modal_data: Optional[MultiModalData] = None,
    ) -> SequenceGroup:
        """Creates a SequenceGroup with PoolingParams."""
        # Defensive copy of PoolingParams, which are used by the pooler
        pooling_params = pooling_params.clone()
        # Create the sequence group.
        seq_group = SequenceGroup(request_id=request_id,
                                  seqs=[seq],
                                  arrival_time=arrival_time,
                                  lora_request=lora_request,
                                  multi_modal_data=multi_modal_data,
                                  pooling_params=pooling_params)
        return seq_group
523

Antoni Baum's avatar
Antoni Baum committed
524
525
    def abort_request(self, request_id: Union[str, Iterable[str]]) -> None:
        """Aborts a request(s) with the given ID.
526
527

        Args:
Antoni Baum's avatar
Antoni Baum committed
528
            request_id: The ID(s) of the request to abort.
529
530
531
532
533
534
535
536
537
538
539

        Details:
            - Refer to the
              :meth:`~vllm.core.scheduler.Scheduler.abort_seq_group`
              from class :class:`~vllm.core.scheduler.Scheduler`.

        Example:
            >>> # initialize engine and add a request with request_id
            >>> request_id = str(0)
            >>> # abort the request
            >>> engine.abort_request(request_id)
540
        """
541
542
        self.scheduler.abort_seq_group(request_id)

543
544
545
546
    def get_model_config(self) -> ModelConfig:
        """Gets the model configuration."""
        return self.model_config

547
548
549
550
    def get_decoding_config(self) -> DecodingConfig:
        """Gets the decoding configuration."""
        return self.decoding_config

551
    def get_num_unfinished_requests(self) -> int:
552
        """Gets the number of unfinished requests."""
553
554
        return self.scheduler.get_num_unfinished_seq_groups()

555
    def has_unfinished_requests(self) -> bool:
556
        """Returns True if there are unfinished requests."""
557
558
        return self.scheduler.has_unfinished_seqs()

559
560
561
562
563
564
565
566
567
568
569
570
    def _process_sequence_group_outputs(
        self,
        seq_group: SequenceGroup,
        outputs: List[EmbeddingSequenceGroupOutput],
    ) -> None:
        seq_group.embeddings = outputs[0].embeddings

        for seq in seq_group.get_seqs():
            seq.status = SequenceStatus.FINISHED_STOPPED

        return

571
    def _process_model_outputs(
572
        self,
573
        output: List[Union[SamplerOutput, PoolerOutput]],
574
        scheduled_seq_groups: List[ScheduledSequenceGroup],
575
576
        ignored_seq_groups: List[SequenceGroup],
        seq_group_metadata_list: List[SequenceGroupMetadata],
577
    ) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
578
        """Apply the model output to the sequences in the scheduled seq groups.
579

580
581
582
        Returns RequestOutputs that can be returned to the client.
        """

583
        now = time.time()
584
585
586
587
588
589

        # Organize outputs by [sequence group][step] instead of
        # [step][sequence group].
        output_by_sequence_group = create_output_by_sequence_group(
            sampler_outputs=output, num_seq_groups=len(scheduled_seq_groups))

590
        # Update the scheduled sequence groups with the model outputs.
591
592
593
        for scheduled_seq_group, outputs, seq_group_meta in zip(
                scheduled_seq_groups, output_by_sequence_group,
                seq_group_metadata_list):
594
            seq_group = scheduled_seq_group.seq_group
595
596
            seq_group.update_num_computed_tokens(
                scheduled_seq_group.token_chunk_size)
597
598
599
            if self.model_config.embedding_mode:
                self._process_sequence_group_outputs(seq_group, outputs)
                continue
600

601
602
            self.output_processor.process_prompt_logprob(seq_group, outputs)
            if seq_group_meta.do_sample:
603
                self.output_processor.process_outputs(seq_group, outputs)
604
605
606

        # Free the finished sequence groups.
        self.scheduler.free_finished_seq_groups()
607
608

        # Create the outputs.
609
610
        request_outputs: List[Union[RequestOutput,
                                    EmbeddingRequestOutput]] = []
611
612
        for scheduled_seq_group in scheduled_seq_groups:
            seq_group = scheduled_seq_group.seq_group
613
            seq_group.maybe_set_first_token_time(now)
614
            request_output = RequestOutputFactory.create(seq_group)
615
            request_outputs.append(request_output)
616
        for seq_group in ignored_seq_groups:
617
            request_output = RequestOutputFactory.create(seq_group)
618
619
620
            request_outputs.append(request_output)
        return request_outputs

621
    def step(self) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
Antoni Baum's avatar
Antoni Baum committed
622
623
        """Performs one decoding iteration and returns newly generated results.

624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
        .. figure:: https://i.imgur.com/sv2HssD.png
            :alt: Overview of the step function
            :align: center

            Overview of the step function.

        Details:
            - Step 1: Schedules the sequences to be executed in the next
              iteration and the token blocks to be swapped in/out/copy.

                - Depending on the scheduling policy,
                  sequences may be `preempted/reordered`.
                - A Sequence Group (SG) refer to a group of sequences
                  that are generated from the same prompt.

639
            - Step 2: Calls the distributed executor to execute the model.
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
            - Step 3: Processes the model output. This mainly includes:

                - Decodes the relevant outputs.
                - Updates the scheduled sequence groups with model outputs
                  based on its `sampling parameters` (`use_beam_search` or not).
                - Frees the finished sequence groups.

            - Finally, it creates and returns the newly generated results.

        Example:
            >>> # Please see the example/ folder for more detailed examples.
            >>>
            >>> # initialize engine and request arguments
            >>> engine = LLMEngine.from_engine_args(engine_args)
            >>> example_inputs = [(0, "What is LLM?",
            >>>    SamplingParams(temperature=0.0))]
            >>>
            >>> # Start the engine with an event loop
            >>> while True:
            >>>     if example_inputs:
            >>>         req_id, prompt, sampling_params = example_inputs.pop(0)
661
            >>>         engine.add_request(str(req_id),prompt,sampling_params)
662
663
664
665
666
667
668
669
670
            >>>
            >>>     # continue the request processing
            >>>     request_outputs = engine.step()
            >>>     for request_output in request_outputs:
            >>>         if request_output.finished:
            >>>             # return or show the request output
            >>>
            >>>     if not (engine.has_unfinished_requests() or example_inputs):
            >>>         break
Antoni Baum's avatar
Antoni Baum committed
671
        """
672
        seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
Antoni Baum's avatar
Antoni Baum committed
673

674
        if not scheduler_outputs.is_empty():
675
            execute_model_req = ExecuteModelRequest(
676
677
678
679
                seq_group_metadata_list=seq_group_metadata_list,
                blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
                blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
                blocks_to_copy=scheduler_outputs.blocks_to_copy,
680
681
682
683
684
                num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
                running_queue_size=scheduler_outputs.running_queue_size,
            )
            output = self.model_executor.execute_model(
                execute_model_req=execute_model_req)
685
686
        else:
            output = []
Antoni Baum's avatar
Antoni Baum committed
687

688
689
        request_outputs = self._process_model_outputs(
            output, scheduler_outputs.scheduled_seq_groups,
690
            scheduler_outputs.ignored_seq_groups, seq_group_metadata_list)
691
692

        # Log stats.
693
        self.do_log_stats(scheduler_outputs, output)
694

695
696
697
698
699
700
701
702
        if not request_outputs:
            # Stop the execute model loop in parallel workers until there are
            # more requests to process. This avoids waiting indefinitely in
            # torch.distributed ops which may otherwise timeout, and unblocks
            # the RPC thread in the workers so that they can process any other
            # queued control plane messages, such as add/remove lora adapters.
            self.model_executor.stop_remote_worker_execution_loop()

703
        return request_outputs
Antoni Baum's avatar
Antoni Baum committed
704

705
706
707
708
    def do_log_stats(
            self,
            scheduler_outputs: Optional[SchedulerOutputs] = None,
            model_output: Optional[List[SamplerOutput]] = None) -> None:
709
710
        """Forced log when no requests active."""
        if self.log_stats:
711
712
            self.stat_logger.log(
                self._get_stats(scheduler_outputs, model_output))
713

714
715
716
717
718
719
720
721
722
723
724
725
    def _get_stats(
            self,
            scheduler_outputs: Optional[SchedulerOutputs],
            model_output: Optional[List[SamplerOutput]] = None) -> Stats:
        """Get Stats to be Logged to Prometheus.

        Args:
            scheduler_outputs: Optional, used to populate metrics related to
                the scheduled batch,
            model_output: Optional, used to emit speculative decoding metrics
                which are created by the workers.
        """
726
        now = time.time()
Woosuk Kwon's avatar
Woosuk Kwon committed
727

728
729
730
731
732
733
734
        # System State
        #   Scheduler State
        num_running_sys = len(self.scheduler.running)
        num_swapped_sys = len(self.scheduler.swapped)
        num_waiting_sys = len(self.scheduler.waiting)

        # KV Cache Usage in %
735
        num_total_gpu = self.cache_config.num_gpu_blocks
736
737
738
739
740
        gpu_cache_usage_sys = 0.
        if num_total_gpu is not None:
            num_free_gpu = self.scheduler.block_manager.get_num_free_gpu_blocks(
            )
            gpu_cache_usage_sys = 1.0 - (num_free_gpu / num_total_gpu)
Woosuk Kwon's avatar
Woosuk Kwon committed
741

742
        num_total_cpu = self.cache_config.num_cpu_blocks
743
        cpu_cache_usage_sys = 0.
744
        if num_total_cpu is not None and num_total_cpu > 0:
745
746
            num_free_cpu = self.scheduler.block_manager.get_num_free_cpu_blocks(
            )
747
748
749
750
751
752
753
            cpu_cache_usage_sys = 1.0 - (num_free_cpu / num_total_cpu)

        # Iteration stats
        num_prompt_tokens_iter = 0
        num_generation_tokens_iter = 0
        time_to_first_tokens_iter: List[float] = []
        time_per_output_tokens_iter: List[float] = []
754
755
        num_preemption_iter = (0 if scheduler_outputs is None else
                               scheduler_outputs.preempted)
756
757
758
759
760
761
762
763
764
765
766
767
768

        # Request stats
        #   Latency
        time_e2e_requests: List[float] = []
        #   Metadata
        num_prompt_tokens_requests: List[int] = []
        num_generation_tokens_requests: List[int] = []
        best_of_requests: List[int] = []
        n_requests: List[int] = []
        finished_reason_requests: List[str] = []

        # NOTE: This loop assumes prefill seq_groups are before
        # decode seq_groups in scheduled_seq_groups.
769
        if scheduler_outputs is not None:
770
            num_generation_tokens_from_prefill_groups = 0.
771
772
773
774
            # NOTE: if scheduler_outputs.num_prefill_groups > 0 and
            # the len of scheduler_outputs.scheduled_seq_groups is !=
            # scheduler_outputs.num_prefill_groups, this means that
            # chunked prefills have been detected.
775
776
777
778

            for idx, scheduled_seq_group in enumerate(
                    scheduler_outputs.scheduled_seq_groups):
                group_was_prefill = idx < scheduler_outputs.num_prefill_groups
779
                seq_group = scheduled_seq_group.seq_group
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807

                # NOTE: a seq_group that completed all of its prefill tokens
                # in the last iteration will have seq_group.is_prefill() = False
                # with group_was_prefill = True
                if group_was_prefill:
                    # Number of prompt tokens.
                    num_prompt_tokens_iter += (
                        scheduled_seq_group.token_chunk_size)

                    # If the seq_group just finished the prefill state
                    # get TTFT.
                    if not seq_group.is_prefill():
                        latency = seq_group.get_last_latency(now)
                        time_to_first_tokens_iter.append(latency)

                        # One generation token per finished prefill.
                        num_generation_tokens_from_prefill_groups += (
                            seq_group.num_seqs())
                else:
                    # TPOTs.
                    latency = seq_group.get_last_latency(now)
                    time_per_output_tokens_iter.append(latency)

                # Because of chunked prefill, we can have a single sequence
                # group that does multiple prompt_runs. To prevent logging
                # the same metadata more than once per request, we standardize
                # on logging request level information for finished requests,
                # which can only happen once.
808
                if seq_group.is_finished():
809
                    # Latency timings
810
811
                    time_e2e_requests.append(now -
                                             seq_group.metrics.arrival_time)
812

813
814
815
816
817
818
819
                    # Metadata
                    num_prompt_tokens_requests.append(
                        len(seq_group.prompt_token_ids))
                    num_generation_tokens_requests.extend([
                        seq.get_output_len()
                        for seq in seq_group.get_finished_seqs()
                    ])
820
821
822
823
                    if seq_group.sampling_params is not None:
                        best_of_requests.append(
                            seq_group.sampling_params.best_of)
                        n_requests.append(seq_group.sampling_params.n)
824
825
826
827
828
829
830
831
832
833
834
835
836
837
                    finished_reason_requests.extend([
                        SequenceStatus.get_finished_reason(seq.status)
                        for seq in seq_group.get_finished_seqs()
                    ])

            # Number of generation tokens.
            #   num_batched_tokens equals the number of prompt_tokens plus the
            #   number of decode_tokens in a single iteration. So,
            #   num_generation_tokens = num_batched_tokens - num_prompt_tokens
            #   + num_generation_tokens_from_prefill_groups (since we generate
            #   one token on prefills on iters where the prefill finishes).
            num_generation_tokens_iter = (
                scheduler_outputs.num_batched_tokens - num_prompt_tokens_iter +
                num_generation_tokens_from_prefill_groups)
838

839
840
841
842
843
844
845
846
        # Spec decode, if enabled, emits specialized metrics from the worker in
        # sampler output.
        if model_output and (model_output[0].spec_decode_worker_metrics
                             is not None):
            spec_decode_metrics = model_output[0].spec_decode_worker_metrics
        else:
            spec_decode_metrics = None

847
848
        return Stats(
            now=now,
849
850
851
852
853
854
855
856
857
858
859
860
861
862
            # System stats
            #   Scheduler State
            num_running_sys=num_running_sys,
            num_swapped_sys=num_swapped_sys,
            num_waiting_sys=num_waiting_sys,
            #   KV Cache Usage in %
            gpu_cache_usage_sys=gpu_cache_usage_sys,
            cpu_cache_usage_sys=cpu_cache_usage_sys,

            # Iteration stats
            num_prompt_tokens_iter=num_prompt_tokens_iter,
            num_generation_tokens_iter=num_generation_tokens_iter,
            time_to_first_tokens_iter=time_to_first_tokens_iter,
            time_per_output_tokens_iter=time_per_output_tokens_iter,
863
            spec_decode_metrics=spec_decode_metrics,
864
            num_preemption_iter=num_preemption_iter,
865
866
867
868
869
870
871
872
873
874

            # Request stats
            #   Latency
            time_e2e_requests=time_e2e_requests,
            #   Metadata
            num_prompt_tokens_requests=num_prompt_tokens_requests,
            num_generation_tokens_requests=num_generation_tokens_requests,
            best_of_requests=best_of_requests,
            n_requests=n_requests,
            finished_reason_requests=finished_reason_requests,
875
876
        )

877
    def add_lora(self, lora_request: LoRARequest) -> bool:
878
        return self.model_executor.add_lora(lora_request)
879
880

    def remove_lora(self, lora_id: int) -> bool:
881
        return self.model_executor.remove_lora(lora_id)
882
883

    def list_loras(self) -> List[int]:
884
        return self.model_executor.list_loras()
885
886

    def check_health(self) -> None:
887
        self.model_executor.check_health()