fix issue and add unit test

Signed-off-by: qizixi <[email protected]>

Author: zixi-qi

tests/v1/core/test_async_scheduler.py

@@ -12,16 +12,20 @@
 
 
 def _make_model_runner_output(
-    scheduler_output: SchedulerOutput, ) -> ModelRunnerOutput:
+        scheduler_output: SchedulerOutput,
+        sampled_token_ids: list[list[int]] | None = None,
+        spec_token_ids: list[list[int]] | None = None) -> ModelRunnerOutput:
     req_ids = list(scheduler_output.num_scheduled_tokens.keys())
+    if not sampled_token_ids:
+        sampled_token_ids = [[i] for i in range(len(req_ids))]
     return ModelRunnerOutput(
         req_ids=req_ids,
         req_id_to_index={
             req_id: i
             for i, req_id in enumerate(req_ids)
         },
-        sampled_token_ids=[[i] for i in range(len(req_ids))],
-        spec_token_ids=None,
+        sampled_token_ids=sampled_token_ids,
+        spec_token_ids=spec_token_ids,
         logprobs=None,
         prompt_logprobs_dict={},
         pooler_output=[],
@@ -55,6 +59,59 @@ def test_stop_by_max_tokens(max_tokens: int):
     assert req1.num_output_tokens == max_tokens
 
 
+def test_spec_decode():
+    max_tokens = 7
+    num_spec_tokens = 3
+    spec_token_ids = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12],
+                      [-1, -2, -3]]
+    sampled_token_ids = [[0], [1, 2, 13], [4, 15], [16], [-1, -2]]
+    scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens,
+                                 async_scheduling=True)
+    requests = create_requests(num_requests=1, max_tokens=max_tokens)
+    req = requests[0]
+
+    sched_outputs: deque[SchedulerOutput] = deque()
+    scheduler.add_request(req)
+    sched_outputs.append(scheduler.schedule())
+    sched_outputs.append(scheduler.schedule())
+
+    i = 0
+    while sched_outputs:
+        sched_output = sched_outputs.popleft()
+        # Overwrite with cached spec decode tokens as done in GPUModelRunner
+        if i > 0:
+            sched_output.scheduled_spec_decode_tokens[
+                req.request_id] = spec_token_ids[i - 1]
+        model_runner_output = _make_model_runner_output(
+            sched_output, [sampled_token_ids[i]], [spec_token_ids[i]])
+        engine_core_output = scheduler.update_from_output(
+            sched_output, model_runner_output)
+        # Validate spec decode stats
+        if engine_core_output:
+            assert engine_core_output[0].scheduler_stats
+            spec_decoding_stats = engine_core_output[
+                0].scheduler_stats.spec_decoding_stats
+            if i == 0:
+                # No spec decode stats for prefill round
+                assert spec_decoding_stats is None
+            else:
+                assert spec_decoding_stats
+                assert spec_decoding_stats.num_drafts == 1
+                assert spec_decoding_stats.num_draft_tokens == num_spec_tokens
+                assert spec_decoding_stats.num_accepted_tokens == len(
+                    sampled_token_ids[i]) - 1
+        sched_output = scheduler.schedule()
+        if sched_output.num_scheduled_tokens:
+            assert sched_output.num_scheduled_tokens[
+                req.request_id] == 1 + num_spec_tokens
+            sched_outputs.append(sched_output)
+        i += 1
+
+    assert scheduler.get_num_unfinished_requests() == 0
+    assert req.num_output_tokens == max_tokens
+    assert req.output_token_ids._x == [0, 1, 2, 13, 4, 15, 16]
+
+
 def test_abort():
     scheduler = create_scheduler(async_scheduling=True)
     requests = create_requests(num_requests=10, max_tokens=20)

vllm/v1/worker/gpu_model_runner.py

@@ -457,6 +457,14 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
                 lora_request=new_req_data.lora_request,
             )
 
+            # Cache computed tokens for new request with
+            # speculative decoding + async scheduling
+            if (self.speculative_config
+                    and self.scheduler_config.async_scheduling):
+                self.cached_num_computed_tokens[req_id] = (
+                    new_req_data.num_computed_tokens +
+                    scheduler_output.num_scheduled_tokens[req_id])
+
             # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
             if self.uses_mrope:
                 image_grid_thw = []
@@ -1754,24 +1762,20 @@ def execute_model(
                 spec_decode_metadata,
                 spec_decode_common_attn_metadata,
             )
+            # Update cached request states for async scheduling
+            if self.scheduler_config.async_scheduling:
+                for idx, req_id in enumerate(self.input_batch.req_ids):
+                    if req_id in self.cached_spec_token_ids:
+                        # Update num computed tokens for running requests
+                        num_rejected_tokens = max_gen_len - len(
+                            valid_sampled_token_ids[idx])
+                        self.cached_num_computed_tokens[
+                            req_id] += scheduler_output.num_scheduled_tokens[
+                                req_id] - num_rejected_tokens
+                    self.cached_spec_token_ids[req_id] = spec_token_ids[idx]
 
         self.eplb_step()
 
-        if self.speculative_config and self.scheduler_config.async_scheduling:
-            assert spec_token_ids
-            for idx, req_id in enumerate(self.input_batch.req_ids):
-                self.cached_spec_token_ids[req_id] = spec_token_ids[idx]
-                num_rejected_tokens = max_gen_len - len(
-                    valid_sampled_token_ids[idx])
-                if req_id not in self.cached_num_computed_tokens:
-                    self.cached_num_computed_tokens[
-                        req_id] = scheduler_output.num_scheduled_tokens[
-                            req_id] - num_rejected_tokens
-                else:
-                    self.cached_num_computed_tokens[
-                        req_id] += scheduler_output.num_scheduled_tokens[
-                            req_id] - num_rejected_tokens
-
         return ModelRunnerOutput(
             req_ids=self.input_batch.req_ids,
             req_id_to_index=self.input_batch.req_id_to_index,