Merge branch 'v0.9.2-dev-fix-pp-preempt' into 'v0.9.2-dev'

fix: 修复pp资源抢占bug

See merge request dcutoolkit/deeplearing/vllm!402

vllm/v1/core/sched/scheduler.py

@@ -281,20 +281,27 @@ class Scheduler(SchedulerInterface):
                     num_draft_tokens=num_draft_tokens,
                     num_lookahead_tokens=self.num_lookahead_tokens)
                 if new_blocks is None:
-                    # The request cannot be scheduled.
-                    # Preempt the lowest-priority request.
-                    if self.policy == SchedulingPolicy.PRIORITY:
-                        preempted_req = max(
-                            self.running,
-                            key=lambda r: (r.priority, r.arrival_time),
-                        )
+                    if new_blocks is None:
+                        if self.use_pp:
+                            preemptable_reqs = [r for r in self.running if
+                                                r.num_tokens_with_spec != r.num_computed_tokens]
+                        else:
+                            preemptable_reqs = self.running
+                        # The request cannot be scheduled.
+                        # Preempt the lowest-priority request.
+                        if self.policy == SchedulingPolicy.PRIORITY:
+                            preempted_req = max(
+                                preemptable_reqs,
+                                key=lambda r: (r.priority, r.arrival_time),
+                            )
+                        else:
+                            preempted_req = preemptable_reqs[-1]
                         self.running.remove(preempted_req)
-                    else:
-                        preempted_req = self.running.pop()
 
                     self.kv_cache_manager.free(preempted_req)
                     preempted_req.status = RequestStatus.PREEMPTED
                     preempted_req.num_computed_tokens = 0
+                    preempted_req.spec_token_ids = []
                     if self.log_stats:
                         preempted_req.record_event(
                             EngineCoreEventType.PREEMPTED, scheduled_timestamp)
@@ -901,20 +908,26 @@ class Scheduler(SchedulerInterface):
                         num_draft_tokens=num_draft_tokens,
                         num_lookahead_tokens=self.num_lookahead_tokens)
                     if new_blocks is None:
+                        if self.use_pp:
+                            preemptable_reqs = [r for r in self.running if
+                                                r.num_tokens_with_spec != r.num_computed_tokens]
+                        else:
+                            preemptable_reqs = self.running
                         # The request cannot be scheduled.
                         # Preempt the lowest-priority request.
                         if self.policy == SchedulingPolicy.PRIORITY:
                             preempted_req = max(
-                                self.running,
+                                preemptable_reqs,
                                 key=lambda r: (r.priority, r.arrival_time),
                             )
-                            self.running.remove(preempted_req)
                         else:
-                            preempted_req = self.running.pop()
+                            preempted_req = preemptable_reqs[-1]
+                        self.running.remove(preempted_req)
 
                         self.kv_cache_manager.free(preempted_req)
                         preempted_req.status = RequestStatus.PREEMPTED
                         preempted_req.num_computed_tokens = 0
+                        preempted_req.spec_token_ids = []
                         if self.log_stats:
                             preempted_req.record_event(
                                 EngineCoreEventType.PREEMPTED, scheduled_timestamp)