[Test] Refactor KL divergence and prefix cache branching to kits (#19715)

python/sglang/test/kits/kl_divergence_kit.py

@@ -0,0 +1,42 @@
+from sglang.test.kl_test_utils import (
+    test_input_output_logprobs_match_decode_cache_hit_helper,
+    test_input_output_logprobs_match_prefill_cache_hit_helper,
+)
+
+
+class KLDivergenceMixin:
+    kl_div_thres: float
+    kl_div_thres_decode: float | None = None
+    kl_div_thres_prefill: float | None = None
+    kl_div_max_samples: int = 32
+    kl_div_prefill_max_new_tokens: int = 512
+    kl_div_decode_max_new_tokens: int = 512
+
+    @classmethod
+    def _build_acc_thresholds(cls, threshold):
+        """Build an ACC_THRESHOLDS dict compatible with kl_test_utils."""
+        return {cls.model: {"kl_div": threshold}}
+
+    @classmethod
+    def test_input_output_logprobs_match_prefill_cache_hit(cls):
+        test_input_output_logprobs_match_prefill_cache_hit_helper(
+            base_url=cls.base_url,
+            ACC_THRESHOLDS=cls._build_acc_thresholds(
+                cls.kl_div_thres_prefill or cls.kl_div_thres
+            ),
+            model_name=cls.model,
+            max_samples=cls.kl_div_max_samples,
+            max_new_tokens=cls.kl_div_prefill_max_new_tokens,
+        )
+
+    @classmethod
+    def test_input_output_logprobs_match_decode_cache_hit(cls):
+        test_input_output_logprobs_match_decode_cache_hit_helper(
+            base_url=cls.base_url,
+            ACC_THRESHOLDS=cls._build_acc_thresholds(
+                cls.kl_div_thres_decode or cls.kl_div_thres
+            ),
+            model_name=cls.model,
+            max_samples=cls.kl_div_max_samples,
+            max_new_tokens=cls.kl_div_decode_max_new_tokens,
+        )

python/sglang/test/kits/prefix_cache_branching_kit.py

@@ -0,0 +1,50 @@
+import requests
+
+
+class PrefixCacheBranchingMixin:
+    cache_chunk_size: int
+
+    @classmethod
+    def send_request_helper(cls, text: str):
+        response = requests.post(
+            cls.base_url + "/generate",
+            json={
+                "text": text,
+                "sampling_params": {
+                    "max_new_tokens": 1,
+                },
+            },
+        )
+        return response.json()
+
+    @classmethod
+    def test_prefix_cache_branching(cls):
+        cls.flush_cache()
+        branching_pos = 257
+        text_prefix = "hi" * branching_pos
+        suffix_list = [
+            "this" * cls.cache_chunk_size * 4,
+            "here" * cls.cache_chunk_size * 4,
+            "that" * cls.cache_chunk_size * 4,
+        ]
+        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 = cls.send_request_helper(text_prefix + suffix)
+            cached_tokens = result["meta_info"]["cached_tokens"]
+            if cache_hit:
+                expected_cached_tokens = (
+                    branching_pos // cls.cache_chunk_size * cls.cache_chunk_size
+                )
+                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"

python/sglang/test/server_fixtures/default_fixture.py

@@ -2,6 +2,8 @@ import logging
 import time
 from contextlib import contextmanager
 
+import requests
+
 from sglang.srt.utils import kill_process_tree
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -65,3 +67,7 @@ class DefaultServerBase(CustomTestCase):
     def tearDownClass(cls):
         kill_process_tree(cls.process.pid)
         time.sleep(2)
+
+    @classmethod
+    def flush_cache(cls):
+        requests.post(cls.base_url + "/flush_cache")