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

[Spec V2] Support specV2 for mamba hybrid attention (#18808)

Browse Source 
Co-authored-by: Yi Zhong <207368749+vincentzed@users.noreply.github.com>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Hanming Lu <hanming@x.ai>
This commit is contained in:
zhangheng
2026-02-27 00:36:01 +08:00
committed by GitHub
parent 78d6674c45
commit e4b708d3e9
8 changed files with 205 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/disaggregation/decode.py
View File

@@ -46,6 +46,7 @@ from sglang.srt.disaggregation.utils import (
poll_and_all_reduce,
prepare_abort,
)
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.managers.schedule_batch import FINISH_ABORT, ScheduleBatch
from sglang.srt.managers.utils import GenerationBatchResult
@@ -169,6 +170,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
speculative_num_draft_tokens: int,
enable_mamba_extra_buffer: bool,
pre_alloc_size: int,
enable_overlap_schedule: bool,
mamba_size: int = None,
):
DecodeReqToTokenPool.__init__(
@@ -179,9 +181,13 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
enable_memory_saver=enable_memory_saver,
pre_alloc_size=pre_alloc_size,
)
self.mamba_ping_pong_track_buffer_size = (
2 if speculative_num_draft_tokens is None else 1
)
if envs.SGLANG_ENABLE_SPEC_V2.get() and not enable_mamba_extra_buffer:
raise ValueError(
"Spec v2 requires mamba scheduler strategy `extra_buffer` for mamba models. "
"Please set `--mamba-scheduler-strategy extra_buffer`."
)
self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1
self.enable_mamba_extra_buffer = enable_mamba_extra_buffer
self.enable_memory_saver = enable_memory_saver
effective_mamba_size = (