chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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 = (

9
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
View File

@@ -170,8 +170,13 @@ class MambaAttnBackendBase(AttentionBackend):
query_start_loc = torch.arange(
0, bs + 1, dtype=torch.int32, device=self.device
)
elif forward_batch.forward_mode.is_extend():
if forward_batch.forward_mode.is_target_verify():
elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
if forward_batch.forward_mode.is_draft_extend_v2():
# HybridLinearAttnBackend.init_forward_metadata calls all sub-backends
# unconditionally, but DRAFT_EXTEND_V2 only runs full-attn layers in
# the draft model, so mamba metadata can be skipped.
query_start_loc = None
elif forward_batch.forward_mode.is_target_verify():
query_start_loc = torch.arange(
0,
forward_batch.input_ids.shape[0] + 1,

11
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -535,7 +535,11 @@ class SchedulerOutputProcessorMixin:
mamba_track_interval = get_global_server_args().mamba_track_interval
if batch.spec_algorithm.is_none() and seq_len % mamba_track_interval == 0:
# for non-spec decode, we update mamba_last_track_seqlen at the end of each track interval
req.mamba_next_track_idx = 1 - req.mamba_next_track_idx
req.mamba_next_track_idx = (
batch.req_to_token_pool.get_mamba_ping_pong_other_idx(
req.mamba_next_track_idx
)
)
req.mamba_last_track_seqlen = seq_len
elif (
not batch.spec_algorithm.is_none()
@@ -548,6 +552,11 @@ class SchedulerOutputProcessorMixin:
!= (actual_seq_len - result.accept_length_per_req_cpu[i])
// mamba_track_interval
):
req.mamba_next_track_idx = (
batch.req_to_token_pool.get_mamba_ping_pong_other_idx(
req.mamba_next_track_idx
)
)
req.mamba_last_track_seqlen = (
actual_seq_len // mamba_track_interval * mamba_track_interval
)

11
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -443,6 +443,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
cache_params: BaseLinearStateParams,
enable_mamba_extra_buffer: bool,
speculative_num_draft_tokens: int = None,
enable_overlap_schedule: bool = True,
):
super().__init__(
size=size,
@@ -450,9 +451,13 @@ class HybridReqToTokenPool(ReqToTokenPool):
device=device,
enable_memory_saver=enable_memory_saver,
)
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
self._init_mamba_pool(

2
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -445,6 +445,7 @@ class ModelRunnerKVCacheMixin:
speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
pre_alloc_size=pre_alloc_size,
enable_overlap_schedule=not self.server_args.disable_overlap_schedule,
mamba_size=self.server_args.max_mamba_cache_size,
)
else:
@@ -468,6 +469,7 @@ class ModelRunnerKVCacheMixin:
cache_params=config.mamba2_cache_params,
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
enable_overlap_schedule=not self.server_args.disable_overlap_schedule,
)
else:
self.req_to_token_pool = ReqToTokenPool(

11
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -232,6 +232,17 @@ class EagleVerifyInputV2Mixin:
device=device,
)
# Set mamba_track_indices for mamba prefix-cache state tracking
if get_global_server_args().enable_mamba_extra_buffer():
batch.mamba_track_indices = torch.tensor(
[
req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx]
for req in batch.reqs
],
dtype=torch.int64,
device=device,
)
# Get a forward batch
batch.forward_mode = (
ForwardMode.IDLE

74
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -785,6 +785,16 @@ class EAGLEWorkerV2(BaseSpecWorker):
accept_index,
) = verify_input.sample(batch, logits_output, vocab_mask)
new_seq_lens = batch.seq_lens + accept_length
# Update mamba state for hybrid GDN models after verification
if (
self.target_worker.model_runner.hybrid_gdn_config is not None
or self.target_worker.model_runner.mamba2_config is not None
):
self._mamba_verify_update(
batch, verify_input, accept_length, accept_index, bs
)
verify_done = torch.get_device_module(self.device).Event()
verify_done.record()
@@ -815,6 +825,70 @@ class EAGLEWorkerV2(BaseSpecWorker):
accept_lens=accept_length,
)
def _mamba_verify_update(
self,
batch: ModelWorkerBatch,
verify_input: EagleVerifyInput,
accept_length: torch.Tensor,
accept_index: torch.Tensor,
bs: int,
):
"""Update mamba state for hybrid GDN models after verification."""
# Calculate accepted_steps for mamba state update
# Include the bonus token (+1)
accepted_length_with_bonus = accept_length
if not batch.forward_mode.is_idle() and accept_index.numel() > 0:
if verify_input.topk != 1:
raise ValueError("Spec v2 currently only supports topk = 1.")
accepted_indices_offset = torch.arange(
0,
bs * self.speculative_num_draft_tokens,
step=self.speculative_num_draft_tokens,
dtype=accepted_length_with_bonus.dtype,
device=accepted_length_with_bonus.device,
)
accepted_steps = accepted_length_with_bonus - 1
if batch.mamba_track_indices is not None:
# If after verify, the request's seq_lens has crossed a mamba track interval,
# we need to update the mamba state for the request at the crossing point.
seq_lens_pre_verify = batch.seq_lens
seq_lens_post_verify = batch.seq_lens + accepted_length_with_bonus
mamba_track_interval = self.server_args.mamba_track_interval
to_track_mask = (
seq_lens_pre_verify // mamba_track_interval
!= seq_lens_post_verify // mamba_track_interval
)
tracking_point = (
seq_lens_post_verify // mamba_track_interval * mamba_track_interval
)
to_track_ith = torch.clamp(
tracking_point - seq_lens_pre_verify - 1, min=0
).to(torch.int64)
req_idx = torch.arange(
bs,
dtype=torch.int64,
device=accepted_length_with_bonus.device,
)
candidate_track_steps = (
accept_index[req_idx, to_track_ith] - accepted_indices_offset
)
mamba_steps_to_track = torch.where(
to_track_mask,
candidate_track_steps,
torch.full_like(candidate_track_steps, -1),
)
else:
mamba_steps_to_track = None
self.target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify(
accepted_steps=accepted_steps,
mamba_track_indices=batch.mamba_track_indices,
mamba_steps_to_track=mamba_steps_to_track,
model=self.target_worker.model_runner.model,
)
def move_accepted_tokens_to_target_kvcache(
self,
batch: ModelWorkerBatch,