chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fa3 & trtllm_mha spec overlap (#11874)

Browse Source
This commit is contained in:
Xinyuan Tong
2025-10-31 17:38:13 -07:00
committed by GitHub
parent 5f98b7fe61
commit 0bfe1d145c
3 changed files with 16 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -90,11 +90,7 @@ class ForwardMode(IntEnum):
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 (include_draft_extend_v2 and self == ForwardMode.DRAFT_EXTEND_V2)
or self == ForwardMode.TARGET_VERIFY
or self == ForwardMode.SPLIT_PREFILL
)
@@ -115,22 +111,21 @@ class ForwardMode(IntEnum):
return self == ForwardMode.TARGET_VERIFY
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
return self == ForwardMode.DRAFT_EXTEND or (
include_v2 and self == ForwardMode.DRAFT_EXTEND_V2
)
def is_draft_extend_v2(self):
# For fixed shape logits output in v2 eagle worker
return self == ForwardMode.DRAFT_EXTEND_V2
def is_extend_or_draft_extend_or_mixed(self):
def is_extend_or_draft_extend_or_mixed(self, include_draft_extend_v2: bool = False):
return (
self == ForwardMode.EXTEND
or self == ForwardMode.DRAFT_EXTEND
or self == ForwardMode.MIXED
or self == ForwardMode.SPLIT_PREFILL
or (include_draft_extend_v2 and self == ForwardMode.DRAFT_EXTEND_V2)
)
def is_cuda_graph(self):