[bugfix] Fix prefill tbo disabled when --deepep-mode=auto (#14333)

Co-authored-by: Cheng Wan <wan4ch@gmail.com>

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -131,6 +131,8 @@ def prepare_mlp_sync_batch_raw(
     ) and not disable_cuda_graph
 
     is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
+    if local_batch is not None:
+        local_batch.is_extend_in_batch = is_extend_in_batch
 
     tbo_preparer = TboDPAttentionPreparer()
     if len(offload_tags) == 0 and disable_overlap_schedule:

test/srt/ep/test_deepep_small.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 from types import SimpleNamespace
 
@@ -224,6 +225,10 @@ class TestTBO(CustomTestCase):
                 "--max-running-requests",
                 "512",
             ],
+            env={
+                **os.environ,
+                "SGLANG_TBO_DEBUG": "1",
+            },
         )
 
     @classmethod
@@ -351,6 +356,10 @@ class TestMTPWithTBO(CustomTestCase):
                 "--max-running-requests",
                 "128",
             ],
+            env={
+                **os.environ,
+                "SGLANG_TBO_DEBUG": "1",
+            },
         )
 
     @classmethod