[Deterministic] Support Qwen3-Next model deterministic inference (#13100)

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):

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),

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()