From 77512ae0d72a7f728f4699f6be6d3ed11f857cd2 Mon Sep 17 00:00:00 2001 From: Yuhao Yao <37280700+yuhyao@users.noreply.github.com> Date: Wed, 3 Dec 2025 17:20:33 +0800 Subject: [PATCH] [bugfix] Fix prefill tbo disabled when --deepep-mode=auto (#14333) Co-authored-by: Cheng Wan --- python/sglang/srt/managers/scheduler_dp_attn_mixin.py | 2 ++ test/srt/ep/test_deepep_small.py | 9 +++++++++ 2 files changed, 11 insertions(+) diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index 16be97703..78ba5dd91 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/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: diff --git a/test/srt/ep/test_deepep_small.py b/test/srt/ep/test_deepep_small.py index 4417affb9..2dd42b189 100644 --- a/test/srt/ep/test_deepep_small.py +++ b/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