[PD]: Support HiCache prefetching and pd-incremental transfer on decode side (#26227)

Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>

test/registered/disaggregation/test_disaggregation_decode_radix_cache.py

@@ -1,11 +1,21 @@
+import asyncio
+import os
+import shutil
+import tempfile
 import time
 import unittest
 from types import SimpleNamespace
 
 import requests
 
+from sglang.benchmark.datasets.random import sample_random_requests
+from sglang.benchmark.utils import get_tokenizer
 from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.kits.cache_hit_kit import run_multiturn_cache_hit_test
+from sglang.test.kits.cache_hit_kit import (
+    async_request_sglang_generate,
+    gen_payload,
+    run_multiturn_cache_hit_test,
+)
 from sglang.test.run_eval import run_eval
 from sglang.test.server_fixtures.disaggregation_fixture import (
     PDDisaggregationServerBase,
@@ -140,5 +150,120 @@ class TestDisaggregationDecodeRadixCacheMooncake(
     transfer_backend_name = "mooncake"
 
 
+@unittest.skipUnless(
+    is_in_ci() or _has_mooncake(),
+    "Mooncake is required for decode radix cache disaggregation coverage.",
+)
+class TestDisaggregationDecodeRadixHiCacheFileBackend(PDDisaggregationServerBase):
+    extra_prefill_args = [
+        "--enable-hierarchical-cache",
+        "--hicache-ratio",
+        "1.2",
+        "--hicache-write-policy",
+        "write_through",
+        "--hicache-storage-backend",
+        "file",
+        "--hicache-storage-prefetch-policy",
+        "wait_complete",
+        "--hicache-io-backend",
+        "kernel",
+        "--hicache-mem-layout",
+        "page_first",
+        "--page-size",
+        "64",
+    ]
+    extra_decode_args = [
+        "--disaggregation-decode-enable-radix-cache",
+        *extra_prefill_args,
+    ]
+    transfer_backend_name = "mooncake"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.hicache_dir = tempfile.mkdtemp(prefix="sglang-hicache-")
+        os.environ["SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR"] = cls.hicache_dir
+
+        super().setUpClass()
+        cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
+        cls.transfer_backend = [
+            "--disaggregation-transfer-backend",
+            cls.transfer_backend_name,
+        ]
+        cls.launch_all()
+
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        os.environ.pop("SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR", None)
+        shutil.rmtree(cls.hicache_dir, ignore_errors=True)
+
+    def _post_ok(self, url):
+        response = requests.post(url, timeout=60)
+        response.raise_for_status()
+
+    def _flush_memory_cache(self):
+        self._post_ok(f"{self.prefill_url}/flush_cache?timeout=30")
+        self._post_ok(f"{self.decode_url}/flush_cache?timeout=30")
+
+    def _generate(self, input_ids, output_len):
+        output = asyncio.run(
+            async_request_sglang_generate(
+                gen_payload(input_ids, output_len),
+                f"{self.base_url}/generate",
+            )
+        )
+        self.assertTrue(output.success, output.error)
+        return output
+
+    def _sample_token_ids(self, input_len, output_len, num_prompts=1):
+        tokenizer = get_tokenizer(self.model)
+        return [
+            list(request.prompt)
+            for request in sample_random_requests(
+                input_len=input_len,
+                output_len=output_len,
+                num_prompts=num_prompts,
+                range_ratio=1.0,
+                tokenizer=tokenizer,
+                dataset_path="",
+                return_text=False,
+            )
+        ]
+
+    def test_decode_hicache_file_backend_l3_reuses_decode_output_after_flush(self):
+        self._post_ok(f"{self.decode_url}/hicache/storage-backend/clear")
+        self._flush_memory_cache()
+
+        num_rounds = 5
+        output_len = 64
+        history = self._sample_token_ids(
+            input_len=256, output_len=output_len, num_prompts=1
+        )[0]
+        suffixes = self._sample_token_ids(
+            input_len=64, output_len=output_len, num_prompts=num_rounds - 1
+        )
+
+        prev_prompt_len = 0
+        prev_output_len = 0
+        for round_idx in range(num_rounds):
+            output = self._generate(history, output_len)
+            if round_idx == 0:
+                self.assertEqual(output.cached_tokens, 0)
+            else:
+                self.assertGreaterEqual(
+                    output.cached_tokens,
+                    prev_prompt_len + prev_output_len,
+                )
+
+            history.extend(output.output_ids)
+            prev_prompt_len = output.prompt_len
+            prev_output_len = len(output.output_ids)
+
+            if round_idx < num_rounds - 1:
+                history.extend(suffixes[round_idx])
+                time.sleep(1)
+                self._flush_memory_cache()
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/managers/test_priority_scheduling_disaggregation.py

@@ -106,11 +106,13 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
         queue._resolve_pending_reqs = MagicMock()
         queue._update_handshake_waiters = MagicMock()
         queue._allocatable_tokens = MagicMock(return_value=1000)
-        queue._pre_alloc = MagicMock(
-            side_effect=lambda req, prefix_indices=None, prefix_len=0: torch.arange(
+
+        def pre_alloc_mock(req, prefix_indices=None, prefix_len=0, total_prefix_len=0):
+            return torch.arange(
                 len(req.origin_input_ids) - prefix_len, dtype=torch.int64
             )
-        )
+
+        queue._pre_alloc = MagicMock(side_effect=pre_alloc_mock)
 
         queue.req_to_token_pool = MagicMock()
         queue.req_to_token_pool.available_size.return_value = 100

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -32,6 +32,7 @@ from unittest.mock import MagicMock
 import torch
 
 from sglang.srt.disaggregation.decode import DecodePreallocQueue
+from sglang.srt.disaggregation.decode_hicache_mixin import DecodePrefixMatch
 from sglang.srt.mem_cache.base_prefix_cache import (
     InsertParams,
     MatchPrefixParams,
@@ -312,7 +313,12 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         queue._resolve_pending_reqs = MagicMock()
         queue._update_handshake_waiters = MagicMock()
         queue._match_prefix_and_lock = MagicMock(
-            return_value=(torch.arange(4, dtype=torch.int64), 4)
+            return_value=DecodePrefixMatch(
+                prefix_indices=torch.arange(4, dtype=torch.int64),
+                l2_host_hit_length=0,
+                l3_storage_hit_length=0,
+                last_device_node=req.last_node,
+            )
         )
         queue._pre_alloc = MagicMock(
             side_effect=AssertionError("_pre_alloc should not run")