[HiCache] Fix SWA L3 cache miss due to a prefetch/hit len mismatch (#27291)

Co-authored-by: Zhangheng <hzh0425@apache.org>

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)

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

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

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."""