From afd2d0b2f46a8f653a5c6b82c135dd2ba748981a Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Mon, 1 Jun 2026 14:43:50 +0800 Subject: [PATCH] [PP][Bugfix] Handle input_ids assignment in prepare_for_extend (#26883) --- python/sglang/srt/managers/scheduler_pp_mixin.py | 7 +++++++ 1 file changed, 7 insertions(+) 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())