diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 350790413..baedd44d5 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -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())