[V1] AsyncLLMEngine

examples/api_client.py

@@ -26,7 +26,6 @@ def post_http_request(prompt: str,
     pload = {
         "prompt": prompt,
         "n": n,
-        "use_beam_search": True,
         "temperature": 0.0,
         "max_tokens": 16,
         "stream": stream,
@@ -58,7 +57,7 @@ def get_response(response: requests.Response) -> List[str]:
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default="localhost")
     parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--n", type=int, default=4)
+    parser.add_argument("--n", type=int, default=1)
     parser.add_argument("--prompt", type=str, default="San Francisco is a")
     parser.add_argument("--stream", action="store_true")
     args = parser.parse_args()
@@ -77,8 +76,8 @@ def get_response(response: requests.Response) -> List[str]:
             num_printed_lines = 0
             for i, line in enumerate(h):
                 num_printed_lines += 1
-                print(f"Beam candidate {i}: {line!r}", flush=True)
+                print(f"Output {i}: {line!r}", flush=True)
     else:
         output = get_response(response)
         for i, line in enumerate(output):
-            print(f"Beam candidate {i}: {line!r}", flush=True)
+            print(f"Output {i}: {line!r}", flush=True)

vllm/entrypoints/api_server.py

@@ -15,15 +15,20 @@
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 
 from vllm.engine.arg_utils import AsyncEngineArgs
-from vllm.engine.async_llm_engine import AsyncLLMEngine
 from vllm.entrypoints.launcher import serve_http
+from vllm.envs import VLLM_USE_V1
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils import (FlexibleArgumentParser, iterate_with_cancellation,
                         random_uuid)
 from vllm.version import __version__ as VLLM_VERSION
 
+if VLLM_USE_V1:
+    from vllm.v1.engine.async_llm_engine import AsyncLLMEngine  # type: ignore
+else:
+    from vllm.engine.async_llm_engine import AsyncLLMEngine  # type: ignore
+
 logger = init_logger("vllm.entrypoints.api_server")
 
 TIMEOUT_KEEP_ALIVE = 5  # seconds.

vllm/entrypoints/llm.py

@@ -1,4 +1,5 @@
 import itertools
+import time
 import warnings
 from contextlib import contextmanager
 from typing import (Any, ClassVar, Dict, List, Optional, Sequence, Tuple,
@@ -14,7 +15,7 @@
                                          apply_hf_chat_template,
                                          apply_mistral_chat_template,
                                          parse_chat_messages)
-from vllm.inputs import PromptType, TextPrompt, TokensPrompt
+from vllm.inputs import INPUT_REGISTRY, PromptType, TextPrompt, TokensPrompt
 from vllm.inputs.parse import parse_and_batch_prompt
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
@@ -32,7 +33,10 @@
 from vllm.utils import Counter, deprecate_args, deprecate_kwargs, is_list_of
 
 if envs.VLLM_USE_V1:
+    from vllm.v1.engine.core import EngineCoreClient  # type: ignore
+    from vllm.v1.engine.detokenizer import Detokenizer  # type: ignore
     from vllm.v1.engine.llm_engine import LLMEngine  # type: ignore
+    from vllm.v1.engine.processor import Processor  # type: ignore
 else:
     from vllm.engine.llm_engine import LLMEngine  # type: ignore
 
@@ -195,8 +199,50 @@
             mm_processor_kwargs=mm_processor_kwargs,
             **kwargs,
         )
-        self.llm_engine = LLMEngine.from_engine_args(
-            engine_args, usage_context=UsageContext.LLM_CLASS)
+
+        # TODO: should this be wrapped in a class?
+        if envs.VLLM_USE_V1:
+            engine_config = engine_args.create_engine_config()
+            executor_class = LLMEngine._get_executor_cls(engine_config)
+            self.task = engine_config.model_config.task
+            self.supported_tasks = engine_config.model_config.task
+
+            # Processor (converts Inputs --> EngineCoreRequests)
+            self.processor = Processor(engine_config.model_config,
+                                       engine_config.parallel_config,
+                                       engine_config.scheduler_config,
+                                       engine_config.lora_config,
+                                       INPUT_REGISTRY)
+
+            # Detokenizer (converts EngineCoreOutputs --> RequestOutput)
+            self.detokenizer = Detokenizer(
+                engine_config.model_config.tokenizer)
+
+            # EngineCoreClient.
+            self.engine_core_client = EngineCoreClient(
+                executor_class,
+                engine_config.model_config,
+                engine_config.cache_config,
+                engine_config.parallel_config,
+                engine_config.scheduler_config,
+                engine_config.device_config,
+                engine_config.load_config,
+                engine_config.lora_config,
+                engine_config.speculative_config,
+                engine_config.decoding_config,
+                engine_config.observability_config,
+                engine_config.prompt_adapter_config,
+                UsageContext.LLM_CLASS,
+                use_async_sockets=False,
+            )
+
+        else:
+            self.llm_engine = LLMEngine.from_engine_args(
+                engine_args, usage_context=UsageContext.LLM_CLASS)
+
+            self.task = self.llm_engine.model_config.task
+            self.supported_tasks = self.llm_engine.model_config.task
+
         self.request_counter = Counter()
 
     def get_tokenizer(self) -> AnyTokenizer:
@@ -337,14 +383,14 @@
             considered legacy and may be deprecated in the future. You should
             instead pass them via the ``inputs`` parameter.
         """
-        task = self.llm_engine.model_config.task
+        task = self.task
         if task != "generate":
             messages = [
                 "LLM.generate() is only supported for (conditional) generation "
                 "models (XForCausalLM, XForConditionalGeneration).",
             ]
 
-            supported_tasks = self.llm_engine.model_config.supported_tasks
+            supported_tasks = self.supported_tasks
             if "generate" in supported_tasks:
                 messages.append(
                     "Your model supports the 'generate' task, but is "
@@ -724,11 +770,11 @@
             considered legacy and may be deprecated in the future. You should
             instead pass them via the ``inputs`` parameter.
         """
-        task = self.llm_engine.model_config.task
+        task = self.task
         if task != "embedding":
             messages = ["LLM.encode() is only supported for embedding models."]
 
-            supported_tasks = self.llm_engine.model_config.supported_tasks
+            supported_tasks = self.supported_tasks
             if "embedding" in supported_tasks:
                 messages.append(
                     "Your model supports the 'embedding' task, but is "
@@ -868,14 +914,28 @@
         priority: int = 0,
     ) -> None:
         request_id = str(next(self.request_counter))
-        self.llm_engine.add_request(
-            request_id,
-            prompt,
-            params,
-            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
-            priority=priority,
-        )
+
+        if envs.VLLM_USE_V1:
+            # 1) Convert input --> DetokenizerRequest / EngineCoreRequest.
+            detokenizer_req, engine_core_req = self.processor.process_inputs(
+                request_id, prompt, params, time.time(), lora_request, None,
+                prompt_adapter_request, priority)
+
+            # 2) Add the request to Detokenizer (this process).
+            self.detokenizer.add_request(detokenizer_req)
+
+            # 3) Add the EngineCoreRequest to EngineCore (separate process).
+            self.engine_core_client.add_request(engine_core_req)
+
+        else:
+            self.llm_engine.add_request(
+                request_id,
+                prompt,
+                params,
+                lora_request=lora_request,
+                prompt_adapter_request=prompt_adapter_request,
+                priority=priority,
+            )
 
     def _add_guided_params(
             self,
@@ -898,9 +958,43 @@
             whitespace_pattern=guided_options.guided_whitespace_pattern)
         return params
 
+    def _run_engine_v1(
+            self, use_tqdm: bool
+    ) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
+        # Initialize tqdm.
+        if use_tqdm:
+            num_requests = self.detokenizer.get_num_unfinished_requests()
+            pbar = tqdm(
+                total=num_requests,
+                desc="Processed prompts",
+                dynamic_ncols=True,
+            )
+
+        # Run the engine.
+        request_outputs: List[Union[RequestOutput,
+                                    EmbeddingRequestOutput]] = []
+        while self.detokenizer.has_unfinished_requests():
+            engine_core_outputs = self.engine_core_client.get_output()
+            outputs = self.detokenizer.step(engine_core_outputs)
+            for output in outputs:
+                if output.finished:
+                    request_outputs.append(output)
+                    if use_tqdm:
+                        pbar.update(1)
+
+        if use_tqdm:
+            pbar.close()
+        # Sort the outputs by request ID.
+        # This is necessary because some requests may be finished earlier than
+        # its previous requests.
+        return sorted(outputs, key=lambda x: int(x.request_id))
+
     def _run_engine(
             self, *, use_tqdm: bool
     ) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
+        if envs.VLLM_USE_V1:
+            return self._run_engine_v1(use_tqdm)
+
         # Initialize tqdm.
         if use_tqdm:
             num_requests = self.llm_engine.get_num_unfinished_requests()
@@ -941,7 +1035,7 @@
         # Sort the outputs by request ID.
         # This is necessary because some requests may be finished earlier than
         # its previous requests.
-        return sorted(outputs, key=lambda x: int(x.request_id))
+        outputs = sorted(outputs, key=lambda x: int(x.request_id))
 
     def _is_encoder_decoder_model(self):
         return self.llm_engine.is_encoder_decoder_model()

vllm/outputs.py

@@ -113,6 +113,36 @@ def __init__(
         self.encoder_prompt = encoder_prompt
         self.encoder_prompt_token_ids = encoder_prompt_token_ids
 
+    @classmethod
+    def create_empty(cls, request_id: str, prompt: Optional[str],
+                     prompt_token_ids: Optional[List[int]]) -> "RequestOutput":
+        """Initialize a new "empty" RequestOutput object."""
+
+        # TODO: Support `n` > 1.
+        completion_output = CompletionOutput(
+            index=0,
+            text="",
+            token_ids=[],
+            cumulative_logprob=None,
+            logprobs=None,  # TODO
+            finish_reason=None,
+            stop_reason=None,
+            lora_request=None,
+        )
+
+        return RequestOutput(
+            request_id=request_id,
+            prompt=prompt,
+            prompt_token_ids=prompt_token_ids,
+            prompt_logprobs=None,  # TODO
+            outputs=[completion_output],
+            finished=False,
+            metrics=None,
+            lora_request=None,
+            encoder_prompt=None,
+            encoder_prompt_token_ids=None,
+        )
+
     @classmethod
     def from_seq_group(
         cls, seq_group: SequenceGroup, use_cache: bool,

vllm/v1/core/scheduler.py

@@ -1,12 +1,13 @@
 from collections import deque
 from dataclasses import dataclass
-from typing import Deque, Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Deque, Dict, Iterable, List, Optional, Set, Union
 
 from vllm.config import CacheConfig, LoRAConfig, SchedulerConfig
 from vllm.logger import init_logger
 from vllm.multimodal import MultiModalDataDict
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.engine import EngineCoreOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
 
@@ -227,13 +228,12 @@ def update_from_output(
         self,
         scheduler_output: "SchedulerOutput",
         model_runner_output: "ModelRunnerOutput",
-    ) -> List[Tuple[Request, int]]:
+    ) -> List[EngineCoreOutput]:
         # NOTE(woosuk): This method doesn't consider speculative decoding.
         sampled_token_ids = model_runner_output.sampled_token_ids_cpu.tolist()
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
         new_running: List[Request] = []
-        # (request, num_sampled_tokens)
-        sampled: List[Tuple[Request, int]] = []
+        engine_core_outputs: List[EngineCoreOutput] = []
         for request in self.running:
             req_id = request.request_id
             request.num_computed_tokens += num_scheduled_tokens[req_id]
@@ -247,17 +247,30 @@ def update_from_output(
                 # generates at most one token at each step.
                 token_id = sampled_token_ids[req_index]
                 request.output_token_ids.append(token_id)
-                sampled.append((request, 1))
+                num_new_tokens = 1
+
                 # TODO: Update the KV cache manager for prefix caching.
 
-                # Check if the request is finished.
+                # Check for stop and update request state.
+                # This must be called before me make the EngineCoreOutput.
                 stopped = self._check_stop(request)
+
+                # Add EngineCoreOutput for this Request.
+                output = EngineCoreOutput(
+                    request_id=req_id,
+                    new_token_ids=request.output_token_ids[-num_new_tokens:],
+                    finished=request.is_finished(),
+                    finish_reason=request.get_finished_reason(),
+                    stop_reason=request.stop_reason)
+                engine_core_outputs.append(output)
+
+                # Breakout of the loop.
                 if stopped:
                     continue
 
             new_running.append(request)
         self.running = new_running
-        return sampled
+        return engine_core_outputs
 
     def _check_stop(self, request: Request) -> bool:
         if (request.num_tokens >= self.max_model_len