diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 5b0c56a84..71e0158a1 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1244,12 +1244,6 @@ class UnifiedRadixCache(BasePrefixCache): self.dec_lock_ref(ancestor_node) return None - logger.info( - "load_back: kv_tokens=%d, node_id=%d", - kv_tokens, - last_hit_node.id, - ) - # Load H→D aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] aux_xfers.extend(anchor_kv_shared_indices_xfers) @@ -1367,14 +1361,9 @@ class UnifiedRadixCache(BasePrefixCache): req = params.req if last_node.evicted or params.host_hit_length > 0: - logger.info( - "init_load_back triggered: node_id=%d, host_hit_length=%d", - last_node.id, - params.host_hit_length, - ) loading_values = self.load_back(last_node, mem_quota, req=req) if loading_values is not None: - logger.info( + logger.debug( "init_load_back success: loaded %d tokens for node %d", len(loading_values), last_node.id, diff --git a/python/sglang/test/kl_test_utils.py b/python/sglang/test/kl_test_utils.py index b3c90caae..de77fb7af 100644 --- a/python/sglang/test/kl_test_utils.py +++ b/python/sglang/test/kl_test_utils.py @@ -118,8 +118,13 @@ def compare_kl_divergence( # Common request helpers -def _flush_cache(base_url): - requests.post(base_url + "/flush_cache") +def _flush_cache(base_url, timeout_s=30): + response = requests.post( + base_url + "/flush_cache", + params={"timeout": timeout_s}, + timeout=timeout_s + 10, + ) + response.raise_for_status() def _generate( diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl.py b/test/registered/radix_cache/test_unified_radix_cache_kl.py index 0543bd5c7..3a58858c4 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl.py @@ -34,7 +34,7 @@ MAMBA_TRACK_INTERVAL = 128 SWA_MODEL = "openai/gpt-oss-20b" FULL_MODEL = "Qwen/Qwen3-32B" -register_cuda_ci(est_time=632, suite="stage-c-test-4-gpu-h100") +register_cuda_ci(est_time=760, suite="stage-c-test-4-gpu-h100") class UnifiedRadixTreeTestMixin: @@ -235,8 +235,8 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): """SWA hybrid + UnifiedRadixCache.""" kl_threshold = 0.03 - gsm8k_threshold = 0.75 - mmlu_threshold = 0.75 + gsm8k_threshold = 0.7 + mmlu_threshold = 0.7 @classmethod def setUpClass(cls): @@ -253,7 +253,7 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): "0.7", "--disable-piecewise-cuda-graph", ], - env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "0"}, + env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, ) cls.input_ids = get_input_ids(cls.model, num_samples=18) @@ -262,54 +262,57 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): kill_process_tree(cls.process.pid) -# TODO(hzh): Currently, HiCache is not fully compatible with the CI CUDA13 environment; we need to wait for the fix before re-enabling the tests below. -# class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): -# """Mamba hybrid + UnifiedRadixCache.""" -# -# kl_threshold = 0.003 -# prefill_cache_assert = staticmethod( -# make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE) -# ) -# decode_cache_assert = staticmethod( -# make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL) -# ) -# -# @classmethod -# def setUpClass(cls): -# cls.model = MAMBA_MODEL -# cls.base_url = DEFAULT_URL_FOR_TEST -# cls.process = popen_launch_server( -# cls.model, -# cls.base_url, -# timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, -# other_args=[ -# "--tp-size", -# "4", -# "--chunked-prefill-size", -# "2048", -# "--mem-fraction-static", -# "0.85", -# "--mamba-scheduler-strategy", -# "extra_buffer", -# "--mamba-track-interval", -# str(MAMBA_TRACK_INTERVAL), -# "--enable-hierarchical-cache", -# "--hicache-ratio", -# "1.5", -# "--hicache-write-policy", -# "write_through", -# "--hicache-io-backend", -# "direct", -# "--hicache-mem-layout", -# "page_first_direct", -# ], -# env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, -# ) -# cls.input_ids = get_input_ids(cls.model, num_samples=18) -# -# @classmethod -# def tearDownClass(cls): -# kill_process_tree(cls.process.pid) +class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): + """Mamba hybrid + HiCache + UnifiedRadixCache.""" + + kl_threshold = 0.003 + prefill_cache_assert = staticmethod( + make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE) + ) + decode_cache_assert = staticmethod( + make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL) + ) + + @classmethod + def setUpClass(cls): + cls.model = MAMBA_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--mem-fraction-static", + "0.85", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + str(MAMBA_TRACK_INTERVAL), + "--enable-hierarchical-cache", + "--hicache-ratio", + "4", + "--hicache-write-policy", + "write_through", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "page_first_direct", + "--max-total-tokens", + "12000", + "--max-running-requests", + "4", + ], + env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + cls.input_ids = get_input_ids(cls.model, num_samples=18) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) if __name__ == "__main__":