[PP][Bugfix] Handle input_ids assignment in prepare_for_extend (#26883)

2026-06-01 14:43:50 +08:00
parent 53b8378307
commit afd2d0b2f4
1 changed files with 7 additions and 0 deletions
@@ -660,6 +660,13 @@ class SchedulerPPMixin:
                start = time.perf_counter()
                batch.prepare_for_extend()
                # Resolve deferred H2D: prepare_for_extend now leaves input_ids=None
                if batch.input_ids is None and batch.prefill_input_ids_cpu is not None:
                    batch.input_ids = batch.prefill_input_ids_cpu.to(
                        self.device, non_blocking=True
                    )
                    batch.prefill_input_ids_cpu = None
                forward_batch = ForwardBatch.init_new(batch, model_runner)
                set_is_extend_in_batch(batch.forward_mode.is_extend())