[core] feat: support async scheduling with structured outputs (#1582)

* [core] feat: support async scheduling with structured outputs

Signed-off-by: AlpinDale <alpindale@gmail.com>

* [cleanup] remove unused variable `num_scheduled_tokens` from GPUModelRunner

Signed-off-by: AlpinDale <alpindale@gmail.com>

---------

Signed-off-by: AlpinDale <alpindale@gmail.com>

Author: AlpinDale

aphrodite/distributed/kv_transfer/kv_connector/utils.py

@@ -4,7 +4,7 @@
 
 from collections.abc import Sequence
 from concurrent.futures import CancelledError, Future
-from typing import TYPE_CHECKING, Literal, cast
+from typing import TYPE_CHECKING, Literal
 
 import torch
 
@@ -138,8 +138,10 @@ def from_connector(cls, connector: "KVConnectorBase", world_size: int):
         return cls(connector.get_finished_count() or world_size)
 
     def aggregate(
-        self, outputs: list[ModelRunnerOutput], output_rank: int = 0
-    ) -> ModelRunnerOutput:
+        self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0
+    ) -> ModelRunnerOutput | None:
+        if not outputs[output_rank]:
+            return None
         # Aggregate kv_connector_output from all workers
 
         def update_finished_set(
@@ -161,6 +163,7 @@ def update_finished_set(
         aggregated_kv_connector_stats = None
         invalid_block_ids = set[int]()
         for model_runner_output in outputs:
+            assert model_runner_output is not None
             kv_output = model_runner_output.kv_connector_output
             if not kv_output:
                 continue
@@ -204,6 +207,7 @@ def update_finished_set(
         # select output of the worker specified by output_rank
         output = outputs[output_rank]
 
+        assert output is not None
         output.kv_connector_output = KVConnectorOutput(
             finished_sending=finished_sending or None,
             finished_recving=finished_recving or None,
@@ -215,13 +219,16 @@ def update_finished_set(
         return output
 
     def async_aggregate(
-        self, output_futures: Sequence[Future[ModelRunnerOutput]], output_rank: int = 0
-    ) -> Future[ModelRunnerOutput]:
+        self,
+        output_futures: Sequence[Future[ModelRunnerOutput | None]],
+        output_rank: int = 0,
+    ) -> Future[ModelRunnerOutput | None]:
         """Takes a list of futures and returns a single future which resolves
         to the respective list of outputs."""
-        result_future: Future[ModelRunnerOutput] = Future()
+        result_future: Future[ModelRunnerOutput | None] = Future()
 
         outputs: list[ModelRunnerOutput | None] = [None] * len(output_futures)
+        remaining = len(output_futures)
 
         def make_callback(idx):
             def callback(fut):
@@ -236,12 +243,10 @@ def callback(fut):
                     result_future.set_exception(e)
 
                 # this check assumes io_thread_pool uses a single thread
-                if all(outputs):
-                    result_future.set_result(
-                        self.aggregate(
-                            cast(list[ModelRunnerOutput], outputs), output_rank
-                        )
-                    )
+                nonlocal remaining
+                remaining -= 1
+                if not remaining:
+                    result_future.set_result(self.aggregate(outputs, output_rank))
 
             return callback
 

aphrodite/v1/core/sched/async_scheduler.py

@@ -12,13 +12,17 @@ def _update_after_schedule(
         scheduler_output: SchedulerOutput,
     ) -> None:
         super()._update_after_schedule(scheduler_output)
+        pending_structured_output_tokens = False
         for req_id in scheduler_output.num_scheduled_tokens:
             request = self.requests[req_id]
+            pending_structured_output_tokens |= request.use_structured_output and request.num_output_placeholders > 0
             if request.num_computed_tokens == request.num_tokens + request.num_output_placeholders:
                 # The request will generate a new token in this scheduling step.
                 # TODO(woosuk): Support speculative decoding.
                 request.num_output_placeholders += 1
 
+        scheduler_output.pending_structured_output_tokens = pending_structured_output_tokens
+
     def _update_request_with_output(
         self,
         request: Request,

aphrodite/v1/core/sched/interface.py

@@ -4,7 +4,7 @@
 
 if TYPE_CHECKING:
     from aphrodite.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
-    from aphrodite.v1.core.sched.output import SchedulerOutput
+    from aphrodite.v1.core.sched.output import GrammarOutput, SchedulerOutput
     from aphrodite.v1.engine import EngineCoreOutputs
     from aphrodite.v1.metrics.stats import SchedulerStats
     from aphrodite.v1.outputs import DraftTokenIds, ModelRunnerOutput
@@ -38,6 +38,10 @@ def schedule(self) -> "SchedulerOutput":
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def get_grammar_bitmask(self, scheduler_output: "SchedulerOutput") -> "GrammarOutput | None":
+        raise NotImplementedError
+
     @abstractmethod
     def update_from_output(
         self,

aphrodite/v1/core/sched/output.py

@@ -173,12 +173,17 @@ class SchedulerOutput:
     # freed from the encoder cache.
     free_encoder_mm_hashes: list[str]
 
-    # ids of structured outputs requests included in the bitmask, in the
-    # same order as the corresponding stacked rows of the bitmask.
-    # There may be more than one row per request in the case of speculative decoding.
-    structured_output_request_ids: list[str]
-    # the bitmask for the whole batch
-    grammar_bitmask: "npt.NDArray[np.int32] | None"
+    # Whether the scheduled requests have all the output tokens they
+    # need to perform grammar bitmask computation.
+    pending_structured_output_tokens: bool = False
 
     # KV Cache Connector metadata.
     kv_connector_metadata: KVConnectorMetadata | None = None
+
+
+@dataclass
+class GrammarOutput:
+    # ids of structured output requests.
+    structured_output_request_ids: list[str]
+    # Bitmask ordered as structured_output_request_ids.
+    grammar_bitmask: "npt.NDArray[np.int32]"

aphrodite/v1/core/sched/scheduler.py

@@ -6,7 +6,7 @@
 import time
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from aphrodite.config import AphroditeConfig
 from aphrodite.distributed.kv_events import EventPublisherFactory, KVEventBatch
@@ -18,7 +18,12 @@
 from aphrodite.v1.core.encoder_cache_manager import EncoderCacheManager, compute_encoder_budget
 from aphrodite.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
 from aphrodite.v1.core.sched.interface import SchedulerInterface
-from aphrodite.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
+from aphrodite.v1.core.sched.output import (
+    CachedRequestData,
+    GrammarOutput,
+    NewRequestData,
+    SchedulerOutput,
+)
 from aphrodite.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from aphrodite.v1.core.sched.utils import check_stop, remove_all
 from aphrodite.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
@@ -29,10 +34,6 @@
 from aphrodite.v1.spec_decode.metrics import SpecDecodingStats
 from aphrodite.v1.structured_output import StructuredOutputManager
 
-if TYPE_CHECKING:
-    import numpy as np
-    import numpy.typing as npt
-
 logger = init_logger(__name__)
 
 
@@ -669,9 +670,6 @@ def schedule(self) -> SchedulerOutput:
             scheduled_spec_decode_tokens,
             req_to_new_blocks,
         )
-        structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
-            num_scheduled_tokens.keys(), scheduled_spec_decode_tokens
-        )
 
         # Record the request ids that were scheduled in this step.
         self.prev_step_scheduled_req_ids.clear()
@@ -691,8 +689,6 @@ def schedule(self) -> SchedulerOutput:
             # the previous and the current steps.
             finished_req_ids=self.finished_req_ids,
             free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
-            structured_output_request_ids=structured_output_request_ids,
-            grammar_bitmask=grammar_bitmask,
         )
 
         # NOTE(Kuntai): this function is designed for multiple purposes:
@@ -907,11 +903,7 @@ def _try_schedule_encoder_inputs(
             encoder_compute_budget,
         )
 
-    def get_grammar_bitmask(
-        self,
-        scheduled_request_ids: Iterable[str],
-        scheduled_spec_decode_tokens: dict[str, list[int]],
-    ) -> tuple[list[str], npt.NDArray[np.int32] | None]:
+    def get_grammar_bitmask(self, scheduler_output: SchedulerOutput) -> GrammarOutput | None:
         # Collect list of scheduled request ids that use structured output.
         # The corresponding rows of the bitmask will be in this order.
         # PERF: in case of chunked prefill,
@@ -920,18 +912,18 @@ def get_grammar_bitmask(
         # cycle to fill in the bitmask, which could be a big no-op.
         structured_output_request_ids = [
             req_id
-            for req_id in scheduled_request_ids
+            for req_id in scheduler_output.num_scheduled_tokens
             if (req := self.requests.get(req_id)) and req.use_structured_output
         ]
         if not structured_output_request_ids:
-            return structured_output_request_ids, None
+            return None
 
         bitmask = self.structured_output_manager.grammar_bitmask(
             self.requests,
             structured_output_request_ids,
-            scheduled_spec_decode_tokens,
+            scheduler_output.scheduled_spec_decode_tokens,
         )
-        return structured_output_request_ids, bitmask
+        return GrammarOutput(structured_output_request_ids, bitmask)
 
     def update_from_output(
         self,

aphrodite/v1/engine/core.py

@@ -10,7 +10,7 @@
 from contextlib import ExitStack, contextmanager
 from inspect import isclass, signature
 from logging import DEBUG
-from typing import Any, TypeVar
+from typing import Any, TypeVar, cast
 
 import msgspec
 import zmq
@@ -168,7 +168,7 @@ def __init__(
         if kv_connector is not None:
             # Collect and store KV connector xfer metadata from workers
             # (after KV cache registration)
-            xfer_handshake_metadata = self.model_executor.get_kv_connector_handshake_metadata()
+            xfer_handshake_metadata = self.modeling.get_kv_connector_handshake_metadata()
 
             if xfer_handshake_metadata:
                 # xfer_handshake_metadata is list of dicts from workers
@@ -305,9 +305,13 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]:
         if not self.scheduler.has_requests():
             return {}, False
         scheduler_output = self.scheduler.schedule()
+        future = self.modeling.execute_model(scheduler_output, non_block=True)
+        grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
 
         with self.log_error_detail(scheduler_output):
-            model_output = self.modeling.execute_model(scheduler_output)
+            model_output = future.result()
+            if model_output is None:
+                model_output = self.modeling.sample_tokens(grammar_output)
 
         engine_core_outputs = self.scheduler.update_from_output(scheduler_output, model_output)
 
@@ -345,16 +349,39 @@ def step_with_batch_queue(
         assert len(batch_queue) < self.batch_queue_size
 
         model_executed = False
+        deferred_scheduler_output = None
         if self.scheduler.has_requests():
             scheduler_output = self.scheduler.schedule()
-            future = self.modeling.execute_model(scheduler_output, non_block=True)
-            batch_queue.appendleft((future, scheduler_output))
+            exec_future = self.modeling.execute_model(scheduler_output, non_block=True)
 
             model_executed = scheduler_output.total_num_scheduled_tokens > 0
-            if model_executed and len(batch_queue) < self.batch_queue_size and not batch_queue[-1][0].done():
-                # Don't block on next worker response unless the queue is full
-                # or there are no more requests to schedule.
-                return None, True
+            if scheduler_output.pending_structured_output_tokens:
+                # We need to defer sampling until we have processed the model output
+                # from the prior step.
+                deferred_scheduler_output = scheduler_output
+                # Block-wait for execute to return (continues running async on the GPU).
+                with self.log_error_detail(scheduler_output):
+                    exec_result = exec_future.result()
+                    assert exec_result is None
+            else:
+                # We aren't waiting for any tokens, get any grammar output immediately.
+                grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
+                # Block-wait for execute to return (continues running async on the GPU).
+                with self.log_error_detail(scheduler_output):
+                    exec_result = exec_future.result()
+
+                if exec_result is None:
+                    # Call sample tokens.
+                    future = self.modeling.sample_tokens(grammar_output, non_block=True)
+                else:
+                    # No sampling required (e.g. all requests finished).
+                    future = cast(Future[ModelRunnerOutput], exec_future)
+                # Add this step's future to the queue.
+                batch_queue.appendleft((future, scheduler_output))
+                if model_executed and len(batch_queue) < self.batch_queue_size and not batch_queue[-1][0].done():
+                    # Don't block on next worker response unless the queue is full
+                    # or there are no more requests to schedule.
+                    return None, True
 
         elif not batch_queue:
             # Queue is empty. We should not reach here since this method should
@@ -368,6 +395,17 @@ def step_with_batch_queue(
             model_output = future.result()
 
         engine_core_outputs = self.scheduler.update_from_output(scheduler_output, model_output)
+
+        # NOTE: We can either handle the deferred tasks here or save
+        # in a field and do it immediately once step_with_batch_queue is
+        # re-called. The latter slightly favors TTFT over TPOT/throughput.
+        if deferred_scheduler_output:
+            # We now have the tokens needed to compute the bitmask for the
+            # deferred request. Get the bitmask and call sample tokens.
+            grammar_output = self.scheduler.get_grammar_bitmask(deferred_scheduler_output)
+            future = self.modeling.sample_tokens(grammar_output, non_block=True)
+            batch_queue.appendleft((future, deferred_scheduler_output))
+
         return engine_core_outputs, model_executed
 
     def shutdown(self):

aphrodite/v1/executor/abstract.py

@@ -12,7 +12,7 @@
 from aphrodite.lora.request import LoRARequest
 from aphrodite.tasks import SupportedTask
 from aphrodite.utils.import_utils import resolve_obj_by_qualname
-from aphrodite.v1.core.sched.output import SchedulerOutput
+from aphrodite.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from aphrodite.v1.engine import ReconfigureDistributedRequest
 from aphrodite.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 from aphrodite.v1.outputs import DraftTokenIds, ModelRunnerOutput
@@ -175,28 +175,44 @@ def get_kv_connector_handshake_metadata(
 
     @overload
     def execute_model(
-        self,
-        scheduler_output: SchedulerOutput,
-        non_block: Literal[False] = False,
-    ) -> ModelRunnerOutput:
+        self, scheduler_output: SchedulerOutput, non_block: Literal[False] = False
+    ) -> ModelRunnerOutput | None:
         pass
 
     @overload
     def execute_model(
-        self,
-        scheduler_output: SchedulerOutput,
-        non_block: Literal[True] = True,
-    ) -> Future[ModelRunnerOutput]:
+        self, scheduler_output: SchedulerOutput, non_block: Literal[True] = True
+    ) -> Future[ModelRunnerOutput | None]:
         pass
 
     def execute_model(
         self, scheduler_output: SchedulerOutput, non_block: bool = False
-    ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
+    ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
         output = self.collective_rpc(  # type: ignore[call-overload]
             "execute_model", args=(scheduler_output,), non_block=non_block
         )
         return output[0]
 
+    @overload
+    def sample_tokens(
+        self, grammar_output: GrammarOutput | None, non_block: Literal[False] = False
+    ) -> ModelRunnerOutput:
+        pass
+
+    @overload
+    def sample_tokens(
+        self, grammar_output: GrammarOutput | None, non_block: Literal[True] = True
+    ) -> Future[ModelRunnerOutput]:
+        pass
+
+    def sample_tokens(
+        self, grammar_output: GrammarOutput | None, non_block: bool = False
+    ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
+        output = self.collective_rpc(  # type: ignore[call-overload]
+            "sample_tokens", args=(grammar_output,), non_block=non_block
+        )
+        return output[0]
+
     def execute_dummy_batch(self) -> None:
         self.collective_rpc("execute_dummy_batch")