[HiCache] refactor: hicache normalization flow and compatibility checks (#19669)

2026-03-20 18:38:44 -07:00
parent 9419453713
commit 6c91590e1b
4 changed files with 184 additions and 46 deletions
@@ -332,5 +332,99 @@ class TestSSLArgs(unittest.TestCase):
        self.assertTrue(server_args.enable_ssl_refresh)


+class TestHiCacheArgs(unittest.TestCase):
+    def _make_args(self, **overrides) -> ServerArgs:
+        args = ServerArgs(model_path="dummy")
+        for key, value in overrides.items():
+            setattr(args, key, value)
+        return args
+
+    def _assert_hicache_fields(
+        self,
+        args: ServerArgs,
+        *,
+        expected_io_backend: str,
+        expected_mem_layout: str,
+        expected_decode_backend: str | None = None,
+    ):
+        self.assertEqual(args.hicache_io_backend, expected_io_backend)
+        self.assertEqual(args.hicache_mem_layout, expected_mem_layout)
+        if expected_decode_backend is not None:
+            self.assertEqual(args.decode_attention_backend, expected_decode_backend)
+
+    def test_hicache_io_backend_and_mem_layout_compatibility(self):
+        cases = [
+            {
+                "name": "kernel_with_page_first_direct",
+                "overrides": {
+                    "enable_hierarchical_cache": True,
+                    "hicache_io_backend": "kernel",
+                    "hicache_mem_layout": "page_first_direct",
+                },
+                "expected_io_backend": "direct",
+                "expected_mem_layout": "page_first_direct",
+            },
+            {
+                "name": "direct_with_page_first",
+                "overrides": {
+                    "enable_hierarchical_cache": True,
+                    "hicache_io_backend": "direct",
+                    "hicache_mem_layout": "page_first",
+                },
+                "expected_io_backend": "direct",
+                "expected_mem_layout": "page_first_direct",
+            },
+            {
+                "name": "mooncake_with_layer_first",
+                "overrides": {
+                    "enable_hierarchical_cache": True,
+                    "hicache_storage_backend": "mooncake",
+                    "hicache_io_backend": "direct",
+                    "hicache_mem_layout": "layer_first",
+                },
+                "expected_io_backend": "direct",
+                "expected_mem_layout": "page_first_direct",
+            },
+            {
+                "name": "fa3_kernel_with_explicit_decode_backend",
+                "overrides": {
+                    "enable_hierarchical_cache": True,
+                    "hicache_io_backend": "kernel",
+                    "hicache_mem_layout": "page_first",
+                    "attention_backend": "triton",
+                    "decode_attention_backend": "fa3",
+                },
+                "expected_io_backend": "direct",
+                "expected_mem_layout": "page_first_direct",
+            },
+        ]
+
+        for case in cases:
+            with self.subTest(case=case["name"]):
+                args = self._make_args(**case["overrides"])
+                args._handle_hicache()
+                self._assert_hicache_fields(
+                    args,
+                    expected_io_backend=case["expected_io_backend"],
+                    expected_mem_layout=case["expected_mem_layout"],
+                )
+
+    @patch.object(ServerArgs, "use_mla_backend", return_value=False)
+    @patch("sglang.srt.server_args.is_flashinfer_available", return_value=False)
+    def test_decode_attention_backend_with_implicit_fa3(
+        self, _mock_flashinfer, _mock_use_mla_backend
+    ):
+        args = self._make_args(
+            enable_hierarchical_cache=True,
+            hicache_io_backend="kernel",
+            attention_backend="fa3",
+            decode_attention_backend=None,
+        )
+
+        args._handle_hicache()
+
+        self.assertEqual(args.decode_attention_backend, "triton")
+
+
 if __name__ == "__main__":
    unittest.main()