[PD] Preserve decode KV across retraction in HiCache (#34801)

Co-authored-by: cctry <cctry@fb.com>
2026-08-17 08:49:11 -07:00
parent af743371cc
commit 2e7c85da68
15 changed files with 779 additions and 34 deletions
@@ -1331,6 +1331,27 @@ class TestHiCacheArgs(unittest.TestCase):
        self.assertEqual(args.hicache_mem_layout, "page_first")
        self.assertIsNone(args.decode_attention_backend)

+    def test_decode_offload_rejects_host_pool_retraction(self):
+        args = self._make_args(
+            disaggregation_mode="decode",
+            disaggregation_decode_enable_offload_kvcache=True,
+            hicache_storage_backend="file",
+            disaggregation_decode_retraction_backup="host_pool",
+        )
+
+        with self.assertRaisesRegex(ValueError, "mutually exclusive"):
+            args._handle_cache_compatibility()
+
+    def test_decode_offload_allows_cpu_tensor_retraction(self):
+        args = self._make_args(
+            disaggregation_mode="decode",
+            disaggregation_decode_enable_offload_kvcache=True,
+            hicache_storage_backend="file",
+            disaggregation_decode_retraction_backup="cpu_tensor",
+        )
+
+        args._handle_cache_compatibility()
+

 class TestNgramExternalSamArgs(CustomTestCase):
    def _make_dummy_ngram_args(self, **overrides):