diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py
index 9c0edc315..5fbab10e5 100644
--- a/python/sglang/srt/managers/scheduler_pp_mixin.py
+++ b/python/sglang/srt/managers/scheduler_pp_mixin.py
@@ -20,6 +20,7 @@ from sglang.srt.layers.dp_attention import (
     get_attention_dp_rank,
     get_attention_dp_size,
     is_dp_attention_enabled,
+    set_is_extend_in_batch,
 )
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.managers.utils import (
@@ -631,6 +632,8 @@ class SchedulerPPMixin:
                 model_worker_batch = batch.get_model_worker_batch()
 
                 forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
+                set_is_extend_in_batch(batch.forward_mode.is_extend())
+
                 _ = model_runner.forward(
                     forward_batch=forward_batch, pp_proxy_tensors=pp_proxy
                 )