[Qwen3-next] support mamba radix cache for overlap scheduler (#14792)

test/srt/models/test_qwen3_next_models.py

@@ -1,6 +1,8 @@
 import unittest
 from types import SimpleNamespace
 
+import requests
+
 from sglang.srt.utils import kill_process_tree
 from sglang.test.few_shot_gsm8k import run_eval
 from sglang.test.kl_test_utils import (
@@ -16,7 +18,22 @@ from sglang.test.test_utils import (
 
 QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
 
-ACC_THRESHOLDS = {QWEN3_NEXT_MODEL: {"kl_div": 0.01, "gsm8k": 0.93}}
+ACC_THRESHOLDS = {
+    QWEN3_NEXT_MODEL: {"kl_div": 0.008, "gsm8k": 0.93},
+}
+
+
+def send_request_helper(base_url: str, text: str):
+    response = requests.post(
+        base_url + "/generate",
+        json={
+            "text": text,
+            "sampling_params": {
+                "max_new_tokens": 1,
+            },
+        },
+    )
+    return response.json()
 
 
 class TestQwen3Next(CustomTestCase):
@@ -33,6 +50,10 @@ class TestQwen3Next(CustomTestCase):
                 "4",
                 "--chunked-prefill-size",
                 "2048",
+                "--mamba-scheduler-strategy",
+                "extra_buffer",
+                "--mamba-track-interval",
+                "128",
             ],
         )
 
@@ -61,8 +82,8 @@ class TestQwen3Next(CustomTestCase):
             self.base_url,
             ACC_THRESHOLDS,
             self.model,
-            max_samples=16,
-            max_new_tokens=256,
+            max_samples=32,
+            max_new_tokens=512,
         )
 
     def test_input_output_logprobs_match_decode_cache_hit(self):
@@ -70,10 +91,37 @@ class TestQwen3Next(CustomTestCase):
             self.base_url,
             ACC_THRESHOLDS,
             self.model,
-            max_samples=16,
-            max_new_tokens=256,
+            max_samples=32,
+            max_new_tokens=512,
         )
 
+    def test_prefix_cache_branching(self):
+        print("running test_prefix_cache_branching")
+        requests.get(self.base_url + "/flush_cache")
+        branching_pos = 257
+        text_prefix = "hi" * branching_pos
+        suffix_list = ["this" * 256, "here" * 256, "that" * 256]
+        cache_hit_list = [False, False, True]
+
+        # First request only prefill the entire sequence
+        # Second request won't have cache hit, but will cache the branching point
+        # Third request will have cache hit on the branching point
+        for i, (suffix, cache_hit) in enumerate(
+            zip(suffix_list, cache_hit_list, strict=True)
+        ):
+            result = send_request_helper(self.base_url, text_prefix + suffix)
+            cached_tokens = result["meta_info"]["cached_tokens"]
+            if cache_hit:
+                expected_cached_tokens = branching_pos // 64 * 64
+                assert (
+                    cached_tokens == expected_cached_tokens
+                ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}"
+            else:
+                assert (
+                    cached_tokens == 0
+                ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0"
+        print("test_prefix_cache_branching passed")
+
 
 class TestQwen3NextMTP(CustomTestCase):
     @classmethod
@@ -98,6 +146,10 @@ class TestQwen3NextMTP(CustomTestCase):
                 "0.8",
                 "--tp",
                 "4",
+                "--chunked-prefill-size",
+                "2048",
+                "--mamba-scheduler-strategy",
+                "no_buffer",
             ],
         )
 
@@ -121,6 +173,24 @@ class TestQwen3NextMTP(CustomTestCase):
             metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"]
         )
 
+    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,
+        )
+
 
 class TestQwen3NextMTPTopk(CustomTestCase):
     @classmethod
@@ -145,6 +215,12 @@ class TestQwen3NextMTPTopk(CustomTestCase):
                 "0.8",
                 "--tp",
                 "4",
+                "--chunked-prefill-size",
+                "2048",
+                "--mamba-scheduler-strategy",
+                "extra_buffer",
+                "--mamba-track-interval",
+                "128",
             ],
         )
 
@@ -168,6 +244,51 @@ class TestQwen3NextMTPTopk(CustomTestCase):
             metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"]
         )
 
+    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,
+        )
+
+    def test_prefix_cache_branching(self):
+        print("running test_prefix_cache_branching")
+        requests.get(self.base_url + "/flush_cache")
+        branching_pos = 257
+        text_prefix = "hi" * branching_pos
+        suffix_list = ["this" * 256, "here" * 256, "that" * 256]
+        cache_hit_list = [False, False, True]
+
+        # First request only prefill the entire sequence
+        # Second request won't have cache hit, but will cache the branching point
+        # Third request will have cache hit on the branching point
+        for i, (suffix, cache_hit) in enumerate(
+            zip(suffix_list, cache_hit_list, strict=True)
+        ):
+            result = send_request_helper(self.base_url, text_prefix + suffix)
+            cached_tokens = result["meta_info"]["cached_tokens"]
+            if cache_hit:
+                expected_cached_tokens = branching_pos // 64 * 64
+                assert (
+                    cached_tokens == expected_cached_tokens
+                ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}"
+            else:
+                assert (
+                    cached_tokens == 0
+                ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0"
+        print("test_prefix_cache_branching passed")
+
 
 class TestQwen3NextPiecewiseCudaGraph(CustomTestCase):
 

test/srt/run_suite.py

@@ -146,7 +146,7 @@ suites = {
         TestFile("test_eagle_dp_attention.py", 200),
     ],
     "per-commit-4-gpu": [
-        TestFile("models/test_qwen3_next_models.py", 472),
+        TestFile("models/test_qwen3_next_models.py", 590),
         TestFile("test_gpt_oss_4gpu.py", 300),
         TestFile("test_local_attn.py", 411),
         TestFile("test_multi_instance_release_memory_occupation.py", 64),

test/srt/test_mamba_unittest.py

@@ -81,10 +81,12 @@ class TestMamba(unittest.TestCase):
         req_to_token_pool = HybridReqToTokenPool(
             size=max_num_reqs,
             mamba_size=mamba_cache_size,
+            mamba_spec_state_size=max_num_reqs,
             max_context_len=max_context_len,
             device=device,
             enable_memory_saver=False,
             cache_params=mamba2_cache_params,
+            enable_mamba_extra_buffer=False,
             speculative_num_draft_tokens=3,
         )
 
@@ -159,10 +161,12 @@ class TestMamba(unittest.TestCase):
         req_to_token_pool = HybridReqToTokenPool(
             size=max_num_reqs,
             mamba_size=mamba_cache_size,
+            mamba_spec_state_size=max_num_reqs,
             max_context_len=max_context_len,
             device=device,
             enable_memory_saver=False,
             cache_params=mamba2_cache_params,
+            enable_mamba_extra_buffer=False,
             speculative_num_draft_tokens=3,
         )
         # setup kv pool