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

Co-authored-by: Cheng Wan <wan4ch@gmail.com>
This commit is contained in:
Yuhao Yao
2025-12-03 17:20:33 +08:00
committed by GitHub
parent c233e9d7a9
commit 77512ae0d7
2 changed files with 11 additions and 0 deletions

View File

@@ -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: