Refactor EngineCoreOutputs

benchmarks/backend_request_func.py

@@ -272,10 +272,15 @@ async def async_request_openai_completions(
         try:
             async with session.post(url=api_url, json=payload,
                                     headers=headers) as response:
+                #print(f"RES = {response.status}")
                 if response.status == 200:
                     first_chunk_received = False
+
+                    #print(response)
+
                     async for chunk_bytes in response.content:
                         chunk_bytes = chunk_bytes.strip()
+                        #print(f"CB = {chunk_bytes}")
                         if not chunk_bytes:
                             continue
 
@@ -313,7 +318,7 @@ async def async_request_openai_completions(
                     else:
                         output.success = False
                         output.error = (
-                            "Never received a valid chunk to calculate TTFT."
+                            "Never received a valid chunk to calculate TTFT. "
                             "This response will be marked as failed!")
                     output.generated_text = generated_text
                     output.latency = most_recent_timestamp - st

benchmarks/benchmark_serving.py

@@ -564,7 +564,10 @@ async def benchmark(
     )
     test_output = await request_func(request_func_input=test_input)
     if not test_output.success:
-        raise ValueError(
+        #raise ValueError(
+        #    "Initial test run failed - Please make sure benchmark arguments "
+        #    f"are correctly specified. Error: {test_output.error}")
+        print(
             "Initial test run failed - Please make sure benchmark arguments "
             f"are correctly specified. Error: {test_output.error}")
     else:

vllm/v1/core/scheduler.py

@@ -11,10 +11,10 @@
 from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
                                                 compute_encoder_budget)
 from vllm.v1.core.kv_cache_manager import KVCacheManager
-from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreOutputs
 from vllm.v1.metrics.stats import SchedulerStats
 from vllm.v1.outputs import ModelRunnerOutput
-from vllm.v1.request import Request, RequestStatus
+from vllm.v1.request import FinishReason, Request, RequestStatus
 
 if TYPE_CHECKING:
     from vllm.multimodal import MultiModalKwargs
@@ -413,11 +413,19 @@ def update_from_output(
         sampled_token_ids = model_runner_output.sampled_token_ids
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
         new_running: List[Request] = []
-        outputs: List[EngineCoreOutput] = []
+        output = EngineCoreOutputs(request_ids=[],
+                                   new_token_id_offsets=[],
+                                   new_token_ids=[],
+                                   finished=[],
+                                   finish_reason={},
+                                   stop_reason=[],
+                                   scheduler_stats=None
+                                   )
 
         # NOTE(woosuk): As len(self.running) can be up to 1K or more, the below
         # loop can be a performance bottleneck. We should do our best to avoid
         # expensive operations inside the loop.
+        offset = 0
         for request in self.running:
             req_id = request.request_id
             num_tokens_scheduled = num_scheduled_tokens.get(req_id, 0)
@@ -455,30 +463,32 @@ def update_from_output(
                 # TODO: Update the KV cache manager for prefix caching.
 
                 # Check for stop and update request state.
-                # This must be called before we make the EngineCoreOutput.
+                # This must be called before we make the EngineCoreOutputs.
                 stopped = self._check_stop(request)
                 if stopped:
                     self._free_request(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)
-                outputs.append(output)
+                # not a list of outputs here
+
+                # Add EngineCoreOutputs for this Request.
+                output.request_ids.append(req_id)
+                output.new_token_id_offsets.append(offset)
+                output.new_token_ids += request.output_token_ids[-num_new_tokens:]
+                output.finished.append(request.is_finished())
+                if request.get_finished_reason() is not None:
+                    output.finish_reason[req_id] = request.get_finished_reason()
+                #print(f"req stop = {request.stop_reason}, {request.status}")
+                output.stop_reason.append(request.stop_reason)
+                offset = offset + 1  # move out of if?
 
                 # Breakout of the loop.
                 if stopped:
                     continue
 
             new_running.append(request)
         self.running = new_running
-        return EngineCoreOutputs(
-            outputs=outputs,
-            scheduler_stats=self.make_stats(),
-        )
+        output.scheduler_stats = self.make_stats()
+        return output
 
     def _check_stop(self, request: Request) -> bool:
         if (request.num_tokens >= self.max_model_len

vllm/v1/engine/__init__.py

@@ -2,7 +2,7 @@
 
 import enum
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 import msgspec
 
@@ -59,17 +59,17 @@ class EngineCoreRequest:
     lora_request: Optional["LoRARequest"]
 
 
-class EngineCoreOutput(
-        msgspec.Struct,
-        array_like=True,  # type: ignore[call-arg]
-        omit_defaults=True,  # type: ignore[call-arg]
-        gc=False):  # type: ignore[call-arg]
-
-    request_id: str
-    new_token_ids: List[int]
-    finished: bool
-    finish_reason: Optional[FinishReason] = None
-    stop_reason: Union[int, str, None] = None
+#class EngineCoreOutput(
+#        msgspec.Struct,
+#        array_like=True,  # type: ignore[call-arg]
+#        omit_defaults=True,  # type: ignore[call-arg]
+#        gc=False):  # type: ignore[call-arg]
+#
+#    request_id: str
+#    new_token_ids: List[int]
+#    finished: bool
+#    finish_reason: Optional[FinishReason] = None
+#    stop_reason: Union[int, str, None] = None
 
 
 class EngineCoreOutputs(
@@ -82,7 +82,12 @@ class EngineCoreOutputs(
     # e.g. columnwise layout
 
     # [num_reqs]
-    outputs: List[EngineCoreOutput]
+    request_ids: List[str]
+    new_token_id_offsets: List[int]
+    new_token_ids: List[int]
+    finished: List[bool]
+    finish_reason: Dict[str, FinishReason]  # Union[List, Dict]?
+    stop_reason: List[Union[int, str, None]]
     scheduler_stats: SchedulerStats
 
 

vllm/v1/engine/async_llm.py

@@ -249,19 +249,27 @@ async def _run_output_handler(self):
                 # Split outputs into chunks of at most
                 # VLLM_V1_OUTPUT_PROC_CHUNK_SIZE, so that we don't block the
                 # event loop for too long.
-                num_outputs = len(outputs.outputs)
+                num_outputs = len(outputs.new_token_id_offsets)
+
                 if num_outputs <= VLLM_V1_OUTPUT_PROC_CHUNK_SIZE:
-                    slices = (outputs.outputs, )
+                    slices = ((0, num_outputs), )
                 else:
-                    slices = np.array_split(
-                        outputs.outputs,
+                    slices = []
+                    parts = np.linspace(
+                        num_outputs,
                         cdiv(num_outputs, VLLM_V1_OUTPUT_PROC_CHUNK_SIZE))
+                    last = 0
+                    for i in parts:
+                        slices.append((last, i))
+                        last = i
+                    print(f"slices = {slices}")
 
                 iteration_stats = None
-                for i, outputs_slice in enumerate(slices):
+                for i, slice in enumerate(slices):
+                    slice_start, slice_end = slice
                     # 2) Process EngineCoreOutputs.
                     processed_outputs = self.output_processor.process_outputs(
-                        outputs_slice, iteration_stats)
+                        outputs, slice_start, slice_end, iteration_stats)
                     # NOTE: RequestOutputs are pushed to their queues.
                     assert not processed_outputs.request_outputs
                     iteration_stats = processed_outputs.iteration_stats

vllm/v1/engine/core.py

@@ -123,7 +123,14 @@ def step(self) -> EngineCoreOutputs:
 
         if not self.scheduler.has_unfinished_requests():
             return EngineCoreOutputs(
-                outputs=[], scheduler_stats=self.scheduler.make_stats())
+                request_ids=[],
+                new_token_id_offsets=[],
+                new_token_ids=[],
+                finished=[],
+                finish_reason={},
+                stop_reason=[],
+                scheduler_stats=self.scheduler.make_stats()
+            )
 
         scheduler_output = self.scheduler.schedule()
         output = self.model_executor.execute_model(scheduler_output)
@@ -299,5 +306,6 @@ def process_output_socket(self, output_path: str):
         with zmq_socket_ctx(output_path, zmq.constants.PUSH) as socket:
             while True:
                 outputs = self.output_queue.get()
+                #print(outputs)
                 encoder.encode_into(outputs, buffer)
                 socket.send_multipart((buffer, ), copy=False)

vllm/v1/engine/detokenizer.py

@@ -8,7 +8,7 @@
 from vllm.sampling_params import RequestOutputKind
 from vllm.transformers_utils.detokenizer_utils import (
     AnyTokenizer, convert_prompt_ids_to_tokens, detokenize_incrementally)
-from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
+from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequest, FinishReason
 
 logger = init_logger(__name__)
 
@@ -98,18 +98,16 @@ def from_new_request(
 
     def update_from_output(
         self,
-        output: EngineCoreOutput,
+        new_token_ids: List[int],
+        finish_reason: Optional[FinishReason],
+        stop_reason: Union[int, str, None],
     ) -> Optional[DetokenizerOutput]:
         """
         Update RequestState for the request_id by:
             1) Detokenize the new token ids incrementally.
             2) Update the RequestOutput with the new text.
         """
 
-        new_token_ids = output.new_token_ids
-        finish_reason = output.finish_reason
-        stop_reason = output.stop_reason
-
         # 1) Detokenize the new token ids incrementally.
         # TODO(woosuk): This method becomes very inefficient when the number of
         # new_token_ids is more than 1. We need to optimize this.

vllm/v1/engine/llm_engine.py

@@ -143,12 +143,14 @@ def add_request(
 
     def step(self) -> List[RequestOutput]:
 
-        # 1) Get EngineCoreOutput from the EngineCore.
+        # 1) Get EngineCoreOutputs from the EngineCore.
         outputs = self.engine_core.get_output()
 
         # 2) Process EngineCoreOutputs.
         processed_outputs = self.output_processor.process_outputs(
-            outputs.outputs)
+            outputs,
+            0,
+            len(outputs.request_ids))
 
         # 3) Abort any reqs that finished due to stop strings.
         self.engine_core.abort_requests(processed_outputs.reqs_to_abort)

vllm/v1/engine/output_processor.py

@@ -7,7 +7,7 @@
 from vllm.outputs import RequestOutput
 from vllm.transformers_utils.detokenizer_utils import AnyTokenizer
 from vllm.transformers_utils.tokenizer_group import BaseTokenizerGroup
-from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest
+from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequest
 from vllm.v1.engine.detokenizer import (DetokenizerOutput,
                                         IncrementalDetokenizer)
 from vllm.v1.metrics.stats import IterationStats, RequestStateStats
@@ -106,59 +106,71 @@ def add_request(
 
     def process_outputs(
         self,
-        engine_core_outputs: List[EngineCoreOutput],
+        engine_core_outputs: EngineCoreOutputs,
+        first: int,
+        last: int,
         iteration_stats: Optional[IterationStats] = None,
     ) -> OutputProcessorOutput:
         """
         Process the EngineCoreOutputs:
         1) Compute stats for logging
         2) Detokenize
         3) Create and handle RequestOutput objects:
-            * If there is a queue (for usage with AsyncLLM), 
+            * If there is a queue (for usage with AsyncLLM),
               put the RequestOutput objects into the queue for
               handling by the per-request generate() tasks.
 
-            * If there is no queue (for usage with LLMEngine), 
+            * If there is no queue (for usage with LLMEngine),
               return a list of RequestOutput objects.
 
         ****************** NOTE FOR DEVELOPERS ******************
 
         VLLM V1 minimizes the number of python loops over the full
-        batch to ensure system overheads are minimized. This is the 
+        batch to ensure system overheads are minimized. This is the
         only function that should loop over EngineCoreOutputs.
 
         If you need to touch every element of the batch, implement a
         method called XXXClass.update_from_output() to be called
         within the loop below. For examples, see:
             * IterationStats.update_from_output()
             * Detokenizer.update_from_output()
-        
+
         TODO(rob): add Protocol makes update_from_output explicit.
-        
+
         **********************************************************
         """
 
         request_outputs: List[RequestOutput] = []
         reqs_to_abort: List[str] = []
         if not iteration_stats:
             iteration_stats = IterationStats(self.log_stats)
-        for engine_core_output in engine_core_outputs:
-            req_id = engine_core_output.request_id
+        for i, req_id in enumerate(engine_core_outputs.request_ids[first:last]):
             req_state = self.request_states.get(req_id)
             if req_state is None:
                 # Ignore output for already-aborted request.
                 continue
 
+            num_tokens = last - first  # might not be robust
+            start = engine_core_outputs.new_token_id_offsets[i]
+            end = engine_core_outputs.new_token_id_offsets[i + 1] if i < num_tokens - 1 else -1
+            # better way to do this?
+            new_token_ids = engine_core_outputs.new_token_ids[start:end]
+
             # 1) Compute stats for this iteration.
-            iteration_stats.update_from_output(engine_core_output,
+            iteration_stats.update_from_output(num_tokens,
                                                req_state.is_prefilling,
                                                req_state.prompt_len,
                                                req_state.stats)
             req_state.is_prefilling = False
 
             # 2) Detokenize the token ids into text.
+            #print(f"finish = {engine_core_outputs.finish_reason.get(req_id)}")
+            #print(f"stop = {engine_core_outputs.stop_reason[i + first]}")
             detokenizer_output = req_state.detokenizer.update_from_output(
-                engine_core_output)
+                new_token_ids,
+                engine_core_outputs.finish_reason.get(req_id),
+                engine_core_outputs.stop_reason[i + first],
+            )
 
             # 3) Create and handle RequestOutput objects.
             if detokenizer_output is not None:
@@ -177,7 +189,7 @@ def process_outputs(
                     assert detokenizer_output.finish_reason is not None
 
                     self.request_states.pop(req_id)
-                    if not engine_core_output.finished:
+                    if not engine_core_outputs.finished[i]:
                         # If req not finished in EngineCore, but Detokenizer
                         # detected stop string, abort needed in EngineCore.
                         reqs_to_abort.append(req_id)