From e4b708d3e9dd90613a9707c5c3cd4e98cd8c4cd8 Mon Sep 17 00:00:00 2001 From: zhangheng Date: Fri, 27 Feb 2026 00:36:01 +0800 Subject: [PATCH] [Spec V2] Support specV2 for mamba hybrid attention (#18808) Co-authored-by: Yi Zhong <207368749+vincentzed@users.noreply.github.com> Co-authored-by: yizhang2077 <1109276519@qq.com> Co-authored-by: Hanming Lu --- python/sglang/srt/disaggregation/decode.py | 12 ++- .../attention/hybrid_linear_attn_backend.py | 9 +- .../scheduler_output_processor_mixin.py | 11 ++- python/sglang/srt/mem_cache/memory_pool.py | 11 ++- .../model_runner_kv_cache_mixin.py | 2 + .../sglang/srt/speculative/eagle_info_v2.py | 11 +++ .../sglang/srt/speculative/eagle_worker_v2.py | 74 ++++++++++++++++ .../test_qwen3_next_models_mtp.py | 84 +++++++++++++++++++ 8 files changed, 205 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 7d8630c5c..0ddccdf46 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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 = ( diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index d700f588b..4fd4d48ac 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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, diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 23f61659b..01eb9a16d 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -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 ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 317be57a3..10b9d9b49 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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( diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 0fccfd027..2af60158f 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -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( diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 8878dfdd3..0ae6e3906 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index f4affc969..03afe03ea 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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, diff --git a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py index 336298ec3..754ea03c4 100644 --- a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py +++ b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import requests +from sglang.srt.environ import envs from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.few_shot_gsm8k import run_eval @@ -211,5 +212,88 @@ class TestQwen3NextMTPTopk(CustomTestCase): print("test_prefix_cache_branching passed") +class TestQwen3NextMTPV2(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN3_NEXT_MODEL + envs.SGLANG_ENABLE_SPEC_V2.set(True) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--mem-fraction-static", + "0.8", + "--tp", + "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + ], + ) + + @classmethod + def tearDownClass(cls): + envs.SGLANG_ENABLE_SPEC_V2.set(False) + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual( + metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + + # TODO(hzh): After merging the PR that fixes specv2 to correctly return log probs, re-open the tests below. https://github.com/sgl-project/sglang/pull/18645 + # def test_input_output_logprobs_match(self): + # test_input_output_logprobs_match_helper( + # self.base_url, + # ACC_THRESHOLDS, + # self.model, + # max_samples=32, + # max_new_tokens=512, + # ) + + # def test_input_output_logprobs_match_prefill_cache_hit(self): + # test_input_output_logprobs_match_prefill_cache_hit_helper( + # self.base_url, + # ACC_THRESHOLDS, + # self.model, + # max_samples=32, + # max_new_tokens=512, + # ) + + # def test_input_output_logprobs_match_decode_cache_hit(self): + # test_input_output_logprobs_match_decode_cache_hit_helper( + # self.base_url, + # ACC_THRESHOLDS, + # self.model, + # max_samples=32, + # max_new_tokens=512, + # ) + + if __name__ == "__main__": unittest.main()