From c461e7714d18fa96d9f162e88e572026b5b2ff53 Mon Sep 17 00:00:00 2001 From: Yineng Zhang Date: Tue, 21 Oct 2025 15:52:15 -0700 Subject: [PATCH] [Auto Sync] Update forward_batch_info.py (20251021) (#11934) Co-authored-by: github-actions[bot] Co-authored-by: yinghui <32845984+cicirori@users.noreply.github.com> --- .../sglang/srt/model_executor/forward_batch_info.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 398bb0da0..03d24cc88 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -84,11 +84,16 @@ class ForwardMode(IntEnum): def is_prefill(self): return self.is_extend() - def is_extend(self): + def is_extend(self, include_draft_extend_v2: bool = False): return ( self == ForwardMode.EXTEND or self == ForwardMode.MIXED or self == ForwardMode.DRAFT_EXTEND + or ( + self == ForwardMode.DRAFT_EXTEND_V2 + if include_draft_extend_v2 + else False + ) or self == ForwardMode.TARGET_VERIFY ) @@ -107,7 +112,11 @@ class ForwardMode(IntEnum): def is_target_verify(self): return self == ForwardMode.TARGET_VERIFY - def is_draft_extend(self): + def is_draft_extend(self, include_v2: bool = False): + if include_v2: + return ( + self == ForwardMode.DRAFT_EXTEND_V2 or self == ForwardMode.DRAFT_EXTEND + ) return self == ForwardMode.DRAFT_EXTEND def is_draft_extend_v2(self):