[test] Add mamba cache release/resume memory test (#14215)

python/sglang/test/test_utils.py

@@ -61,6 +61,8 @@ DEFAULT_MLA_FP8_MODEL_NAME_FOR_TEST = "neuralmagic/DeepSeek-Coder-V2-Lite-Instru
 DEFAULT_MODEL_NAME_FOR_TEST_MLA = "lmsys/sglang-ci-dsv3-test"
 DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN = "lmsys/sglang-ci-dsv3-test-NextN"
 
+# Hybrid Mamba models
+DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST = "Qwen/Qwen3-Next-80B-A3B-Instruct"
 # VL test models
 DEFAULT_MODEL_NAME_FOR_TEST_VL_PP = "Qwen/Qwen3-VL-2B-Thinking"
 DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP = "zai-org/GLM-4.1V-9B-Thinking"

test/srt/test_release_memory_occupation.py

@@ -39,6 +39,7 @@ from sglang.srt.constants import (
     GPU_MEMORY_TYPE_WEIGHTS,
 )
 from sglang.test.test_utils import (
+    DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST,
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
     DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
@@ -90,6 +91,10 @@ class TestReleaseMemoryOccupation(CustomTestCase):
             "sampling_params_moe": {"temperature": 0, "max_new_tokens": 16},
             "expect_output_before_update_weights_moe": " go to the park. I have a picnic basket, a book, and a",
             "expect_output_after_update_weights_moe": " go to the park. I have a lot of things to do, but I",
+            "prompt_hybrid_mamba": "The weather is nice today, and I want to",
+            "sampling_params_hybrid_mamba": {"temperature": 0, "max_new_tokens": 16},
+            "expect_output_before_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
+            "expect_output_after_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
         }
 
     def _test_initial_generation(
@@ -400,6 +405,66 @@ class TestReleaseMemoryOccupation(CustomTestCase):
         self.assertEqual(outputs, params["expect_output_after_update_weights_moe"])
         engine.shutdown()
 
+    def test_hybrid_mamba_model_release_and_resume(self):
+        # Test with Hybrid Mamba model
+        model_name = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
+
+        tp_size = 4
+
+        print(
+            f"Testing tp_size={tp_size} for test_hybrid_mamba_model_release_and_resume"
+        )
+        engine = sgl.Engine(
+            model_path=model_name,
+            random_seed=42,
+            enable_memory_saver=True,
+            tp_size=tp_size,
+        )
+        params = self._common_test_params()
+
+        self._test_initial_generation(
+            engine,
+            params["prompt_hybrid_mamba"],
+            params["sampling_params_hybrid_mamba"],
+            params["expect_output_before_update_weights_hybrid_mamba"],
+        )
+
+        t = time.perf_counter()
+        gpu_memory_usage_before_release = get_gpu_memory_gb()
+        engine.release_memory_occupation()
+        gpu_memory_usage_after_release = get_gpu_memory_gb()
+        self.assertLess(
+            gpu_memory_usage_after_release,
+            gpu_memory_usage_before_release,
+        )
+
+        print(
+            f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
+        )
+
+        if _DEBUG_EXTRA:
+            time.sleep(3)
+
+        t = time.perf_counter()
+        engine.resume_memory_occupation()
+        print(
+            f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
+        )
+
+        engine.update_weights_from_disk(model_name)
+
+        # destroy the hf model
+        torch.cuda.empty_cache()
+
+        print("generate (#2)")
+        outputs = engine.generate(
+            params["prompt_hybrid_mamba"], params["sampling_params_hybrid_mamba"]
+        )["text"]
+        self.assertEqual(
+            outputs, params["expect_output_after_update_weights_hybrid_mamba"]
+        )
+        engine.shutdown()
+
 
 if __name__ == "__main__":
     unittest.main()