diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1057626b5..6d546ee41 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2201,7 +2201,10 @@ class Scheduler( if last_host_node.backuped or last_host_node is self.tree_cache.root_node: last_hash = last_host_node.get_last_hash_value() matched_len = len(req.prefix_indices) + req.host_hit_length - new_input_tokens = req.full_untruncated_fill_ids[matched_len:] + match_end = req._compute_max_prefix_len( + len(req.full_untruncated_fill_ids) + ) + new_input_tokens = req.full_untruncated_fill_ids[matched_len:match_end] prefix_keys = ( last_host_node.get_prefix_hash_values(last_host_node.parent) diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index 617defda6..2fb5ce7bd 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py @@ -113,6 +113,9 @@ class TestUnifiedDeepSeekV4FlashHiCachePageFirstDirect( class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): """DeepSeek V4 Flash FP8 + HiCache L3 (file backend) + UnifiedRadixCache.""" + l3_prefetch_page_size = 256 + l3_prefetch_prompt_pages = 4 + @classmethod def setUpClass(cls): cls.model = DSV4_FLASH_MODEL @@ -169,6 +172,8 @@ class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestC """DeepSeek V4 Flash EAGLE + HiCache L3 should load from storage.""" page_size = 256 + l3_prefetch_page_size = 256 + l3_prefetch_prompt_pages = 4 input_ids = list(range(4000, 4300)) storage_wait_timeout = 120 num_gsm8k_questions = 100 diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py index c11b91029..969a9653b 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py @@ -25,6 +25,7 @@ register_cuda_ci(est_time=768, stage="base-c", runner_config="4-gpu-h100") MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" MAMBA_CHUNK_SIZE = 64 MAMBA_TRACK_INTERVAL = 128 +MAMBA_CHUNKED_PREFILL_SIZE = 2048 class TestUnifiedMambaRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): @@ -51,7 +52,7 @@ class TestUnifiedMambaRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): "--tp-size", "4", "--chunked-prefill-size", - "2048", + str(MAMBA_CHUNKED_PREFILL_SIZE), "--mem-fraction-static", "0.85", "--mamba-scheduler-strategy", @@ -94,7 +95,7 @@ class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): "--tp-size", "4", "--chunked-prefill-size", - "2048", + str(MAMBA_CHUNKED_PREFILL_SIZE), "--mem-fraction-static", "0.85", "--mamba-scheduler-strategy", @@ -133,6 +134,13 @@ class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): """Mamba hybrid + HiCache L3 (file backend) + UnifiedRadixCache.""" + # Prompt must exceed chunked_prefill_size to exercise the multi-chunk path. + l3_prefetch_page_size = MAMBA_CHUNK_SIZE + l3_prefetch_prompt_pages = MAMBA_CHUNKED_PREFILL_SIZE // MAMBA_CHUNK_SIZE + 16 + # Mamba state is only persisted at chunk boundaries, so up to a full + # chunked_prefill_size of trailing tokens may stay uncached. + l3_prefetch_max_uncached_tokens = MAMBA_CHUNKED_PREFILL_SIZE + @classmethod def setUpClass(cls): cls.model = MAMBA_MODEL @@ -146,7 +154,7 @@ class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): "--tp-size", "4", "--chunked-prefill-size", - "2048", + str(MAMBA_CHUNKED_PREFILL_SIZE), "--mem-fraction-static", "0.85", "--mamba-scheduler-strategy", diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py index 4a2b0953d..0b8302439 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py @@ -6,6 +6,7 @@ via KL divergence. """ import os +import random import shutil import tempfile import unittest @@ -42,6 +43,12 @@ class AccuracyTwoPassMixin: max_accuracy_diff: float = 0.02 + l3_prefetch_page_size: int = 64 + l3_prefetch_prompt_pages: int = 16 + # Max tokens that may stay uncached on a full-prompt re-request; the bound + # depends on model architecture. Defaults to page_size; subclasses override. + l3_prefetch_max_uncached_tokens: int = None + def _run_gsm8k(self): from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k @@ -106,6 +113,32 @@ class AccuracyTwoPassMixin: """Run GSM8K twice with flush in between, verify accuracy diff <= max_accuracy_diff.""" self._two_pass("GSM8K", self._run_gsm8k, self.gsm8k_threshold) + def test_l3_prefetch_full_prefix_hit_after_flush(self): + from sglang.test.kl_test_utils import _flush_cache, _generate + + page = int(self.l3_prefetch_page_size) + n_tokens = page * int(self.l3_prefetch_prompt_pages) + max_uncached = int( + self.l3_prefetch_max_uncached_tokens + if self.l3_prefetch_max_uncached_tokens is not None + else page + ) + + rng = random.Random(987) + input_ids = [rng.randint(1, 30000) for _ in range(n_tokens)] + + _generate(self.base_url, [input_ids], max_new_tokens=4) + _flush_cache(self.base_url) + results = _generate(self.base_url, [input_ids], max_new_tokens=4) + cached = int(results[0]["meta_info"]["cached_tokens"]) + + expected_min = n_tokens - max_uncached + self.assertGreaterEqual( + cached, + expected_min, + f"cached_tokens={cached} < {expected_min} (= input_len - {max_uncached})", + ) + class TestGLM5HiRadixCacheL3Accuracy(AccuracyTwoPassMixin, CustomTestCase): """GLM-5.1-FP8 + HiCache L3 (file backend), with HiRadixTree."""