Merge branch 'v0.9.2-dev-fix-pp+cp' into 'v0.9.2-dev'

V0.9.2 dev fix pp+cp

See merge request dcutoolkit/deeplearing/vllm!364

vllm/v1/core/sched/scheduler.py

@@ -1101,13 +1101,14 @@ class Scheduler(SchedulerInterface):
             req_id = req.request_id
             req_ids.append(req_id)
             num_tokens = req.num_generated_token_ids
+
             if self.use_pp:
                 # When using PP, the scheduler sends the sampled tokens back,
                 # because there's no direct communication between the first-
                 # stage worker and the last-stage worker. Otherwise, we don't
                 # need to send the sampled tokens back because the model runner
                 # will cache them.
-                token_ids = req.all_token_ids[-num_tokens:]
+                token_ids = req.all_token_ids[-num_tokens:] if num_tokens > 0 else []
                 new_token_ids.append(token_ids)
             new_block_ids.append(req_to_new_block_ids[req_id])
             num_computed_tokens.append(req.num_computed_tokens)
@@ -1241,7 +1242,7 @@ class Scheduler(SchedulerInterface):
 
             scheduled_spec_token_ids = (
                 scheduler_output.scheduled_spec_decode_tokens.get(req_id))
-            request.num_generated_token_ids = 1
+            request.num_generated_token_ids = len(generated_token_ids)
             if scheduled_spec_token_ids:
                 # num_computed_tokens represents the number of tokens
                 # processed in the current step, considering scheduled
@@ -1253,7 +1254,6 @@ class Scheduler(SchedulerInterface):
                 num_tokens_rejected = (len(scheduled_spec_token_ids) + 1 -
                                        len(generated_token_ids))
                 request.num_computed_tokens -= num_tokens_rejected
-                request.num_generated_token_ids = len(generated_token_ids)
                 spec_decoding_stats = self.make_spec_decoding_stats(
                     spec_decoding_stats,
                     num_draft_tokens=len(scheduled_spec_token_ids),

vllm/v1/worker/gpu_model_runner.py

@@ -512,14 +512,10 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
                 new_token_ids = req_data.new_token_ids[i]
                 # Add the sampled token(s) from the previous step (if any).
                 # This doesn't include "unverified" tokens like spec tokens.
-                num_new_tokens = (num_computed_tokens + len(new_token_ids) -
-                                  req_state.num_tokens)
-                if num_new_tokens == 1:
-                    # Avoid slicing list in most common case.
-                    req_state.output_token_ids.append(new_token_ids[-1])
-                elif num_new_tokens > 0:
+                num_new_tokens = len(new_token_ids)
+                if num_new_tokens > 0:
                     req_state.output_token_ids.extend(
-                        new_token_ids[-num_new_tokens:])
+                        new_token_ids)
             if len(spec_token_ids) > 0:
                 req_state.spec_token_ids = spec_token_ids
 
@@ -539,6 +535,11 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
                 # The request is not in the persistent batch.
                 # The request was either preempted and resumed later, or was not
                 # scheduled in the previous step and needs to be added again.
+                if not is_last_rank:
+                    req_state = self.requests[req_id]
+                    self.input_batch.add_request(req_state)
+                    req_index = self.input_batch.req_id_to_index.get(req_id)
+                else:
                     req_ids_to_add.append(req_id)
                     continue
 
@@ -552,6 +553,7 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
             if not is_last_rank:
                 # Add new_token_ids to token_ids_cpu.
                 start_token_index = num_computed_tokens
+                if len(new_token_ids) > 0:
                     end_token_index = num_computed_tokens + 1
                     self.input_batch.token_ids_cpu[
                         req_index,
@@ -1643,9 +1645,6 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
         # Mask out the sampled tokens that should not be sampled.
         for i in discard_sampled_tokens_req_indices:
             valid_sampled_token_ids[i].clear()
-        if spec_token_ids is not None:
-            for i in discard_sampled_tokens_req_indices:
-                spec_token_ids[i].clear()
 
         # Cache the sampled tokens in the model runner, so that the scheduler
         # doesn't need to send them back.
@@ -1685,7 +1684,9 @@ class GPUModelRunnerBase(LoRAModelRunnerMixin):
                 spec_decode_metadata,
                 attn_metadata,
             )
-
+        if spec_token_ids is not None:
+            for i in discard_sampled_tokens_req_indices:
+                spec_token_ids[i].clear()
         # Clear KVConnector state after all KVs are generated.
         if has_kv_transfer_group():
             get_kv_transfer_group().clear_connector_metadata()
@@ -3492,7 +3493,6 @@ class GPUModelRunnerMTP(GPUModelRunnerBase):
             hidden_states[:num_scheduled_tokens],
             scheduler_output,
         )
-        #-----------------------------------
         # Get the valid generated tokens.
         sampled_token_ids = sampler_output.sampled_token_ids
         max_gen_len = sampled_token_ids.shape[-1]