diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 739e28943..c89f5f6d0 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -297,6 +297,15 @@ class MambaPool: self.mamba_cache.temporal[:, free_index] = 0 def clear(self): + # Zero the entire mamba cache before resetting free_slots + # This ensures that when slots are reallocated, they start with clean state + if self.is_kda_cache: + for i in range(len(self.mamba_cache.conv)): + self.mamba_cache.conv[i].zero_() + else: + self.mamba_cache.conv.zero_() + self.mamba_cache.temporal.zero_() + self.free_slots = torch.arange(self.size, dtype=torch.int64, device=self.device) def copy_from(self, src_index: torch.Tensor, dst_index: torch.Tensor): diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 51175ab66..359688931 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -223,7 +223,9 @@ suites = { TestFile("nightly/test_flashinfer_trtllm_gen_attn_backend.py", 300), ], "nightly-8-gpu-b200": [], - "nightly-4-gpu": [], + "nightly-4-gpu": [ + TestFile("test_qwen3_next_deterministic.py", 200), + ], "nightly-8-gpu": [], "nightly-8-gpu-h200": [ TestFile("test_deepseek_v32_nsabackend.py", 600), diff --git a/test/srt/test_qwen3_next_deterministic.py b/test/srt/test_qwen3_next_deterministic.py new file mode 100644 index 000000000..70ae4cb9e --- /dev/null +++ b/test/srt/test_qwen3_next_deterministic.py @@ -0,0 +1,87 @@ +""" +Usage: +cd test/srt +python3 -m unittest test_qwen3_next_deterministic.TestFlashInferDeterministic +""" + +import unittest + +from sglang.test.test_deterministic_utils import ( + COMMON_SERVER_ARGS, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + TestDeterministicBase, + popen_launch_server, +) + +QWEN3_NEXT = "Qwen/Qwen3-Next-80B-A3B-Instruct" + + +class TestFlashInferDeterministic(TestDeterministicBase): + @classmethod + def get_model(cls): + return QWEN3_NEXT + + # Test with flashinfer attention backend + @classmethod + def get_server_args(cls): + args = COMMON_SERVER_ARGS + args.extend(["--attention-backend", "flashinfer", "--tp", "4"]) + return args + + @classmethod + def setUpClass(cls): + import os + + cls.model = cls.get_model() + cls.base_url = DEFAULT_URL_FOR_TEST + if "--attention-backend" not in cls.get_server_args(): + raise unittest.SkipTest("Skip the base test class") + + custom_env = os.environ.copy() + custom_env["SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM"] = "false" + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=cls.get_server_args(), + env=custom_env, + ) + + +class TestTritonDeterministic(TestDeterministicBase): + @classmethod + def get_model(cls): + return QWEN3_NEXT + + # Test with triton attention backend + @classmethod + def get_server_args(cls): + args = COMMON_SERVER_ARGS + args.extend(["--attention-backend", "triton", "--tp", "4"]) + return args + + @classmethod + def setUpClass(cls): + import os + + cls.model = cls.get_model() + cls.base_url = DEFAULT_URL_FOR_TEST + if "--attention-backend" not in cls.get_server_args(): + raise unittest.SkipTest("Skip the base test class") + + custom_env = os.environ.copy() + custom_env["SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM"] = "false" + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=cls.get_server_args(), + env=custom_env, + ) + + +if __name__ == "__main__": + unittest.main()