diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index f462f3301..83aeed081 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -144,6 +144,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): enable_memory_saver: bool, cache_params: "Mamba2CacheParams", speculative_num_draft_tokens: int, + enable_mamba_extra_buffer: bool, pre_alloc_size: int, ): DecodeReqToTokenPool.__init__( @@ -154,10 +155,11 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): enable_memory_saver=enable_memory_saver, pre_alloc_size=pre_alloc_size, ) - self.enable_memory_saver = enable_memory_saver - self.enable_mamba_extra_buffer = ( - False # TODO: add PD support for mamba cache extra_buffer + self.mamba_ping_pong_track_buffer_size = ( + 2 if speculative_num_draft_tokens is None else 1 ) + self.enable_mamba_extra_buffer = enable_mamba_extra_buffer + self.enable_memory_saver = enable_memory_saver self._init_mamba_pool( size=size + pre_alloc_size, mamba_spec_state_size=size + pre_alloc_size, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index a7e267b74..0395e8407 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1797,6 +1797,7 @@ class ModelRunner: enable_memory_saver=self.server_args.enable_memory_saver, cache_params=config.mamba2_cache_params, 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, ) else: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 75f7e4762..5abd4f418 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1382,9 +1382,6 @@ class ServerArgs: assert ( is_cuda() ), "Mamba extra_buffer is only supported on CUDA devices with FLA backend" - assert ( - self.disaggregation_mode == "null" - ), "Mamba extra_buffer is not compatible with disaggregation mode yet." if self.speculative_num_draft_tokens is not None: assert ( self.mamba_track_interval >= self.speculative_num_draft_tokens diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index c90a4e5b8..87647998a 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -157,7 +157,7 @@ suites = { "per-commit-8-gpu-h200": [ TestFile("test_deepseek_v3_basic.py", 275), TestFile("test_deepseek_v3_mtp.py", 275), - TestFile("test_disaggregation_hybrid_attention.py", 200), + TestFile("test_disaggregation_hybrid_attention.py", 600), TestFile("models/test_kimi_k2_models.py", 200), TestFile("test_deepseek_v32_basic.py", 275), TestFile("test_deepseek_v32_mtp.py", 275), diff --git a/test/srt/test_disaggregation_hybrid_attention.py b/test/srt/test_disaggregation_hybrid_attention.py index 792a6df05..cf310515c 100644 --- a/test/srt/test_disaggregation_hybrid_attention.py +++ b/test/srt/test_disaggregation_hybrid_attention.py @@ -79,5 +79,77 @@ class TestDisaggregationHybridAttentionMamba(PDDisaggregationServerBase): self.assertGreater(metrics["accuracy"], 0.93) +class TestDisaggregationHybridAttentionMambaExtraBuffer(PDDisaggregationServerBase): + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct" + + # Non blocking start servers + cls.start_prefill() + cls.start_decode() + + # Block until both + cls.wait_server_ready(cls.prefill_url + "/health") + cls.wait_server_ready(cls.decode_url + "/health") + + cls.launch_lb() + + @classmethod + def start_prefill(cls): + prefill_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "prefill", + "--tp", + "4", + "--mamba-scheduler-strategy", + "extra_buffer", + ] + prefill_args += cls.transfer_backend + cls.rdma_devices + cls.process_prefill = popen_launch_pd_server( + cls.model, + cls.prefill_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=prefill_args, + ) + + @classmethod + def start_decode(cls): + decode_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "decode", + "--tp", + "4", + "--base-gpu-id", + "4", + "--mamba-scheduler-strategy", + "extra_buffer", + ] + decode_args += cls.transfer_backend + cls.rdma_devices + cls.process_decode = popen_launch_pd_server( + cls.model, + cls.decode_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=decode_args, + ) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host=f"http://{self.base_host}", + port=int(self.lb_port), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"Evaluation metrics: {metrics}") + + self.assertGreater(metrics["accuracy"], 0.93) + + if __name__ == "__main__": unittest.main()