[Perf] Async Scheduling + Speculative Decoding + Structured Outputs

vllm/v1/core/sched/interface.py

@@ -88,8 +88,22 @@ def update_from_output(
     def update_draft_token_ids(
         self,
         draft_token_ids: "DraftTokenIds",
+        update_requests: bool = True,
+        update_scheduler_output: "SchedulerOutput | None" = None,
+        pad_filtered_draft_tokens: bool = False,
     ) -> None:
-        """Update the draft token ids for the scheduled requests."""
+        """Update requests with newly generated draft token ids, applying
+        structured output grammar validation if needed.
+
+        Args:
+            draft_token_ids: The input draft token ids for each request.
+            update_requests: If True, update the scheduler's `self.requests` state
+            with the draft token ids.
+            update_scheduler_output: If provided, update the given scheduler_output
+            with the corresponding draft token ids.
+            pad_filtered_draft_tokens: If True, pad the draft token ids so that
+            the length does not change after filtering invalid draft tokens.
+        """
         raise NotImplementedError
 
     @abstractmethod

vllm/v1/core/sched/output.py

@@ -228,3 +228,5 @@ class GrammarOutput:
     structured_output_request_ids: list[str]
     # Bitmask ordered as structured_output_request_ids.
     grammar_bitmask: "npt.NDArray[np.int32]"
+    # Number of invalid tokens per structured output request.
+    num_invalid_tokens_per_req: list[int] | None = None

vllm/v1/core/sched/scheduler.py

@@ -1281,6 +1281,9 @@ def _free_encoder_inputs(self, request: Request) -> None:
     def update_draft_token_ids(
         self,
         draft_token_ids: DraftTokenIds,
+        update_requests: bool = True,
+        update_scheduler_output: SchedulerOutput | None = None,
+        pad_filtered_draft_tokens: bool = False,
     ) -> None:
         for req_id, spec_token_ids in zip(
             draft_token_ids.req_ids,
@@ -1291,14 +1294,25 @@ def update_draft_token_ids(
                 # The request may have been finished. Skip.
                 continue
 
-            # Add newly generated spec token ids to the request.
+            # Filter out spec tokens which do not adhere to the grammar.
+            orig_num_spec_tokens = len(spec_token_ids)
             if self.structured_output_manager.should_advance(request):
                 metadata = request.structured_output_request
-                request.spec_token_ids = metadata.grammar.validate_tokens(  # type: ignore[union-attr]
-                    spec_token_ids
+                assert metadata is not None and metadata.grammar is not None
+                spec_token_ids = metadata.grammar.validate_tokens(
+                    spec_token_ids,
                 )
-            else:
+                if pad_filtered_draft_tokens:
+                    num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)
+                    spec_token_ids.extend([-1] * num_invalid_tokens)
+
+            # Update the scheduler state.
+            if update_requests:
                 request.spec_token_ids = spec_token_ids
+            if update_scheduler_output is not None:
+                update_scheduler_output.scheduled_spec_decode_tokens[req_id] = (
+                    spec_token_ids
+                )
 
     def get_request_counts(self) -> tuple[int, int]:
         """Returns (num_running_reqs, num_waiting_reqs)."""

vllm/v1/engine/core.py

@@ -448,11 +448,39 @@ def step_with_batch_queue(
         # 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.
+            # If we are doing speculative decoding with structured output,
+            # we need to get the draft token ids from the prior step before
+            # we can compute the grammar bitmask for the deferred request.
+            draft_token_ids = self.model_executor.take_draft_token_ids()
+            num_invalid_spec_tokens = None
+            if draft_token_ids is not None:
+                # Update the draft token ids on the scheduler output
+                # to filter out the invalid spec tokens, which will be padded with -1
+                # and ignored by the grammar bitmask computation.
+                self.scheduler.update_draft_token_ids(
+                    draft_token_ids,
+                    update_requests=False,
+                    update_scheduler_output=deferred_scheduler_output,
+                    pad_filtered_draft_tokens=True,
+                )
+                scheduled_spec_tokens = (
+                    deferred_scheduler_output.scheduled_spec_decode_tokens
+                )
+                num_invalid_spec_tokens = {
+                    req_id: sum(token_id == -1 for token_id in spec_token_ids)
+                    for req_id, spec_token_ids in scheduled_spec_tokens.items()
+                }
+
+            # Compute the grammar bitmask using the draft tokens (if any),
+            # and then unblock the model executor.
             grammar_output = self.scheduler.get_grammar_bitmask(
                 deferred_scheduler_output
             )
+            if num_invalid_spec_tokens and grammar_output is not None:
+                grammar_output.num_invalid_tokens_per_req = [
+                    num_invalid_spec_tokens.get(req_id, 0)
+                    for req_id in grammar_output.structured_output_request_ids
+                ]
             future = self.model_executor.sample_tokens(grammar_output, non_block=True)
             batch_queue.appendleft((future, deferred_scheduler_output))
 

vllm/v1/engine/input_processor.py

@@ -157,12 +157,11 @@ def _validate_supported_sampling_params(
                 or params.presence_penalty != 0.0
                 or params.repetition_penalty != 1.0
                 or params.bad_words_token_ids
-                or params.structured_outputs
             )
         ):
             raise ValueError(
                 "async scheduling with spec decoding doesn't yet support "
-                "penalties, bad words or structured outputs in sampling parameters."
+                "penalties or bad words in sampling parameters."
             )
 
     def _validate_params(

vllm/v1/structured_output/__init__.py

@@ -263,7 +263,9 @@ def grammar_bitmask(
                             )
                         ]
                     )
-
+                    if token == -1:
+                        # Stop advancing the grammar once we hit a padding token
+                        apply_bitmask = False
                     if (
                         apply_bitmask
                         and token is not None

vllm/v1/worker/gpu_model_runner.py

@@ -574,6 +574,7 @@ def __init__(
 
         # Cached outputs.
         self._draft_token_ids: list[list[int]] | torch.Tensor | None = None
+        self._draft_token_req_ids: list[str] | None = None
         self.transfer_event = torch.Event()
         self.sampled_token_ids_pinned_cpu = torch.empty(
             (self.max_num_reqs, 1),
@@ -582,6 +583,10 @@ def __init__(
             pin_memory=self.pin_memory,
         )
 
+        self.invalid_spec_tokens_mask = self._make_buffer(
+            (self.max_num_reqs, 1 + self.num_spec_tokens), dtype=torch.bool
+        )
+
         # Pre-allocated tensor for copying valid sampled token counts to CPU,
         # with dedicated stream for overlapping and event for coordination.
         self.valid_sampled_token_count_event: torch.Event | None = None
@@ -596,6 +601,20 @@ def __init__(
             pin_memory=self.pin_memory,
         )
 
+        # We also copy the drafted tokens to the CPU asynchronously,
+        # in case we need them for structured outputs.
+        self.draft_token_ids_event: torch.Event | None = None
+        self.draft_token_ids_copy_stream: torch.cuda.Stream | None = None
+        if self.use_async_scheduling:
+            self.draft_token_ids_event = torch.Event()
+            self.draft_token_ids_copy_stream = torch.cuda.Stream()
+        self.draft_token_ids_cpu = torch.empty(
+            (self.max_num_reqs, self.num_spec_tokens),
+            dtype=torch.int64,
+            device="cpu",
+            pin_memory=self.pin_memory,
+        )
+
         # Ephemeral state transferred between execute_model() and sample_tokens().
         self.execute_model_state: ExecuteModelState | None = None
         self.kv_connector_output: KVConnectorOutput | None = None
@@ -1231,6 +1250,7 @@ def _prepare_input_ids(
         )
 
         # Scatter the draft tokens after the sampled tokens are scattered.
+        self._prev_draft_token_ids = self._draft_token_ids
         if self._draft_token_ids is None or not spec_flattened_indices:
             return
 
@@ -1245,7 +1265,6 @@ def _prepare_input_ids(
         # because input_ids dtype is torch.int32,
         # so convert draft_token_ids to torch.int32 here.
         draft_token_ids = self._draft_token_ids.to(dtype=torch.int32)
-        self._draft_token_ids = None
 
         self.input_ids.gpu.scatter_(
             dim=0,
@@ -3089,10 +3108,13 @@ def execute_model(
 
     @torch.inference_mode
     def sample_tokens(
-        self, grammar_output: "GrammarOutput | None"
+        self,
+        grammar_output: "GrammarOutput | None",
     ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
         kv_connector_output = self.kv_connector_output
         self.kv_connector_output = None
+        self._draft_token_ids = None
+        self._draft_token_req_ids = None
 
         if self.execute_model_state is None:
             # Nothing to do (PP non-final rank case), output isn't used.
@@ -3132,6 +3154,29 @@ def sample_tokens(
         with record_function_or_nullcontext("gpu_model_runner: sample"):
             sampler_output = self._sample(logits, spec_decode_metadata)
 
+        # Mask out invalid spec tokens for async scheduling + structured outputs.
+        if (
+            grammar_output is not None
+            and grammar_output.num_invalid_tokens_per_req is not None
+            and self.use_async_scheduling
+        ):
+            num_invalid_tokens_per_req = grammar_output.num_invalid_tokens_per_req
+            num_reqs, num_sampled_toks = sampler_output.sampled_token_ids.shape
+            num_invalid_spec_tokens_cpu = torch.zeros((num_reqs,), dtype=torch.int32)
+            for req_id, num_invalid_toks in zip(
+                grammar_output.structured_output_request_ids, num_invalid_tokens_per_req
+            ):
+                req_index = self.input_batch.req_id_to_index[req_id]
+                num_invalid_spec_tokens_cpu[req_index] = num_invalid_toks
+            col_indices = torch.arange(num_sampled_toks, dtype=torch.int32)
+            mask_start_indices = num_sampled_toks - num_invalid_spec_tokens_cpu
+            mask = col_indices.unsqueeze(0) >= mask_start_indices.unsqueeze(1)
+            self.invalid_spec_tokens_mask.cpu[:num_reqs, :].copy_(mask)
+            self.invalid_spec_tokens_mask.copy_to_gpu(num_reqs)
+            sampler_output.sampled_token_ids.masked_fill_(
+                self.invalid_spec_tokens_mask.gpu[:num_reqs, :], -1
+            )
+
         self.input_batch.prev_sampled_token_ids = None
 
         def propose_draft_token_ids(sampled_token_ids):
@@ -3147,6 +3192,8 @@ def propose_draft_token_ids(sampled_token_ids):
                     spec_decode_metadata,
                     spec_decode_common_attn_metadata,
                 )
+                self._copy_draft_token_ids_to_cpu(self._draft_token_ids)
+                self._draft_token_req_ids = self.input_batch.req_ids.copy()
 
         spec_config = self.speculative_config
         use_padded_batch_for_eagle = (
@@ -3265,14 +3312,38 @@ def propose_draft_token_ids(sampled_token_ids):
     def take_draft_token_ids(self) -> DraftTokenIds | None:
         if self._draft_token_ids is None:
             return None
-        req_ids = self.input_batch.req_ids
-        if isinstance(self._draft_token_ids, torch.Tensor):
-            draft_token_ids = self._draft_token_ids.tolist()
-        else:
-            draft_token_ids = self._draft_token_ids
-        self._draft_token_ids = None
+        req_ids = self._draft_token_req_ids
+        assert req_ids is not None
+        draft_token_ids = self._get_draft_token_ids_cpu(len(req_ids))
         return DraftTokenIds(req_ids, draft_token_ids)
 
+    def _copy_draft_token_ids_to_cpu(
+        self, draft_token_ids: torch.Tensor | list[list[int]]
+    ) -> None:
+        if isinstance(draft_token_ids, list):
+            return
+        if self.draft_token_ids_event is None:
+            return
+        # For async scheduling, trigger async copy of draft token ids to cpu.
+        default_stream = torch.cuda.current_stream()
+        with torch.cuda.stream(self.draft_token_ids_copy_stream):
+            assert self.draft_token_ids_copy_stream is not None
+            self.draft_token_ids_copy_stream.wait_stream(default_stream)
+            self.draft_token_ids_cpu[: draft_token_ids.shape[0]].copy_(
+                draft_token_ids, non_blocking=True
+            )
+            self.draft_token_ids_event.record()
+
+    def _get_draft_token_ids_cpu(self, num_reqs: int) -> list[list[int]]:
+        if isinstance(self._draft_token_ids, list):
+            return self._draft_token_ids
+        if self.draft_token_ids_event is None:
+            if self._draft_token_ids is None:
+                return []
+            return self._draft_token_ids.tolist()
+        self.draft_token_ids_event.synchronize()
+        return self.draft_token_ids_cpu[0:num_reqs].tolist()
+
     def _copy_valid_sampled_token_count(
         self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor
     ) -> None:

vllm/v1/worker/gpu_worker.py

@@ -535,7 +535,8 @@ def annotate_profile(self, scheduler_output):
 
     @torch.inference_mode()
     def sample_tokens(
-        self, grammar_output: "GrammarOutput | None"
+        self,
+        grammar_output: "GrammarOutput | None",
     ) -> ModelRunnerOutput | AsyncModelRunnerOutput:
         return self.model_runner.sample_tokens(grammar_output)